diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 0c616a59..e20d9289 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -26,7 +26,10 @@ concurrency: cancel-in-progress: true jobs: pytest: - name: pytest + name: pytest (${{ matrix.mink }}) + strategy: + matrix: + mink: ["mink==1.2.0", "mink"] runs-on: ubuntu-latest timeout-minutes: 10 permissions: @@ -37,6 +40,7 @@ jobs: - name: Sync environment run: | uv sync + uv pip install "${{ matrix.mink }}" - name: Run pytest run: | - uv run pytest + uv run --no-sync pytest diff --git a/pyproject.toml b/pyproject.toml index efe08904..64fc14b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,9 +26,9 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ - "daqp==0.7.2", - "mink>=1.1.0", - "mujoco>=3.6.0", + "daqp>=0.8.2", + "mink>=1.2.0", + "mujoco>=3.8.1", "numpy", "openarm-mujoco>=2.0.1", "pyyaml", diff --git a/src/openarm_control/geometry/jacobian.py b/src/openarm_control/geometry/jacobian.py index 1f21065e..8ada8b1e 100644 --- a/src/openarm_control/geometry/jacobian.py +++ b/src/openarm_control/geometry/jacobian.py @@ -22,6 +22,32 @@ import numpy.typing as npt +class CachedConfiguration(mink.Configuration): + """Cache read-only geometric quantities until the next configuration update.""" + + def __init__(self, model: mujoco.MjModel, q: np.ndarray | None = None) -> None: + """Initialize caches before Mink's constructor calls update.""" + self._frame_jacobians: dict[tuple[str | int, str], np.ndarray] = {} + self._arm_svds: dict[mink.Task, tuple[np.ndarray, np.ndarray]] = {} + super().__init__(model, q=q) + + def update(self, q: np.ndarray | None = None) -> None: + """Invalidate geometry whenever forward kinematics is updated.""" + self._frame_jacobians.clear() + self._arm_svds.clear() + super().update(q=q) + + def get_frame_jacobian(self, frame_name: str | int, frame_type: str) -> np.ndarray: + """Return a read-only body Jacobian shared by tasks at this configuration.""" + key = (frame_name, frame_type) + jacobian = self._frame_jacobians.get(key) + if jacobian is None: + jacobian = super().get_frame_jacobian(frame_name, frame_type) + jacobian.setflags(write=False) + self._frame_jacobians[key] = jacobian + return jacobian + + def relative_root_is_independent_of_dofs( frame_task: mink.Task, model: mujoco.MjModel, @@ -34,25 +60,39 @@ def relative_root_is_independent_of_dofs( if not isinstance(native_task, mink.RelativeFrameTask): raise TypeError("Expected a Mink frame task or a frame-task wrapper.") - root_id = mujoco.mj_name2id( - model, - { - "body": mujoco.mjtObj.mjOBJ_BODY, - "site": mujoco.mjtObj.mjOBJ_SITE, - "geom": mujoco.mjtObj.mjOBJ_GEOM, - }[native_task.root_type], - native_task.root_name, + return frame_is_independent_of_dofs( + model, native_task.root_name, native_task.root_type, dof_indices ) - if root_id < 0: - raise ValueError( - f"Unknown {native_task.root_type} frame {native_task.root_name!r}." + + +def frame_is_independent_of_dofs( + model: mujoco.MjModel, + frame_name: str | int, + frame_type: str, + dof_indices: npt.ArrayLike, +) -> bool: + """Return whether none of the selected DoFs is an ancestor of the frame.""" + frame_id = ( + mujoco.mj_name2id( + model, + { + "body": mujoco.mjtObj.mjOBJ_BODY, + "site": mujoco.mjtObj.mjOBJ_SITE, + "geom": mujoco.mjtObj.mjOBJ_GEOM, + }[frame_type], + frame_name, ) - if native_task.root_type == "body": - body_id = root_id - elif native_task.root_type == "site": - body_id = int(model.site_bodyid[root_id]) + if isinstance(frame_name, str) + else int(frame_name) + ) + if frame_id < 0: + raise ValueError(f"Unknown {frame_type} frame {frame_name!r}.") + if frame_type == "body": + body_id = frame_id + elif frame_type == "site": + body_id = int(model.site_bodyid[frame_id]) else: - body_id = int(model.geom_bodyid[root_id]) + body_id = int(model.geom_bodyid[frame_id]) selected_dofs = set(np.asarray(dof_indices, dtype=int).tolist()) while body_id > 0: diff --git a/src/openarm_control/kinematics.py b/src/openarm_control/kinematics.py index d8bc6070..5c0149e3 100644 --- a/src/openarm_control/kinematics.py +++ b/src/openarm_control/kinematics.py @@ -40,6 +40,7 @@ import numpy as np from openarm_control.config import ArmSetup, frame_name +from openarm_control.geometry.jacobian import CachedConfiguration from openarm_control.geometry.poses import pose_to_se3 from openarm_control.ik_params import ( IKParams, @@ -155,7 +156,7 @@ def __init__(self, setup: ArmSetup, params: IKParams) -> None: self._substep_dt = params.dt / params.max_iters self._max_iters = params.max_iters - self._config = mink.Configuration(setup.model) + self._config = CachedConfiguration(setup.model) self._config.update(q=setup.data.qpos.copy()) mid_qpos = self._config.data.qpos.copy() @@ -271,6 +272,13 @@ def __init__(self, setup: ArmSetup, params: IKParams) -> None: self._singularity_limits[side] = limit self._limits.append(limit) + self._paired_singularity_limits: ( + tuple[SingularityApproachLimit, SingularityApproachLimit] | None + ) = None + limits = tuple(self._singularity_limits.values()) + if len(limits) == 2 and limits[0].can_share_perturbations(limits[1]): + self._paired_singularity_limits = (limits[0], limits[1]) + self._posture_task = mink.PostureTask(setup.model, cost=params.posture_cost) self._posture_task.set_target(mid_qpos) @@ -291,6 +299,10 @@ def __init__(self, setup: ArmSetup, params: IKParams) -> None: characteristic_length=params.jacobian_characteristic_length, ) + nullspace_tasks = tuple(self._nullspace_tasks.values()) + if len(nullspace_tasks) == 2: + nullspace_tasks[0]._batch_svd_with(nullspace_tasks[1]) + self._kinetic_energy_task: KineticEnergyRegularizationTask | None = None if params.kinetic_energy_cost > 0.0: self._kinetic_energy_task = KineticEnergyRegularizationTask( @@ -357,8 +369,22 @@ def solve(self) -> np.ndarray | None: constraints = [self._freeze_task] if self._freeze_task else [] q_before = self._config.data.qpos.copy() - for limit in self._singularity_limits.values(): - limit.prepare(self._config) + gradients: tuple[np.ndarray | None, ...] = (None,) * len( + self._singularity_limits + ) + if self._paired_singularity_limits is not None: + first, second = self._paired_singularity_limits + gradients = first._paired_gradients(self._config, second) + for (side, limit), gradient in zip( + self._singularity_limits.items(), gradients, strict=True + ): + nullspace = self._nullspace_tasks.get(side) + command_ratio = ( + nullspace.compute_singularity_ratio(self._config) + if nullspace is not None + else None + ) + limit.prepare(self._config, command_ratio=command_ratio, gradient=gradient) self._pending = set(self._sides) for _ in range(self._max_iters): diff --git a/src/openarm_control/qp/arm_joint_limit.py b/src/openarm_control/qp/arm_joint_limit.py index f5418065..e1c55b94 100644 --- a/src/openarm_control/qp/arm_joint_limit.py +++ b/src/openarm_control/qp/arm_joint_limit.py @@ -24,8 +24,8 @@ import numpy.typing as npt -class ArmConfigurationLimit(mink.ConfigurationLimit): - """Apply Mink's configuration limits only to selected scalar arm joints.""" +class ArmConfigurationLimit(mink.Limit): + """Apply configuration limits only to selected scalar arm joints.""" def __init__( self, @@ -35,8 +35,14 @@ def __init__( gain: float, ) -> None: """Build a configuration limit projected onto selected arm DoFs.""" - super().__init__(model, gain=gain) + if not 0.0 < gain <= 1.0: + raise mink.exceptions.LimitDefinitionError( + f"{self.__class__.__name__} gain must be in the range (0, 1]" + ) + self.model = model + self.gain = gain selected_qpos = {int(index) for index in qpos_indices} + active_joints: list[int] = [] active_dofs: list[int] = [] for joint_id in range(model.njnt): if ( @@ -49,13 +55,34 @@ def __init__( int(mujoco.mjtJoint.mjJNT_SLIDE), ): raise ValueError("ArmConfigurationLimit only supports scalar joints.") + active_joints.append(joint_id) active_dofs.append(int(model.jnt_dofadr[joint_id])) self.indices = _readonly(active_dofs, dtype=int) + self.qpos_indices = _readonly(model.jnt_qposadr[active_joints], dtype=int) + self.lower = _readonly(model.jnt_range[active_joints, 0]) + self.upper = _readonly(model.jnt_range[active_joints, 1]) self.projection_matrix = ( np.eye(model.nv, dtype=np.float64)[self.indices] if self.indices.size else None ) + self._G = ( + np.vstack([self.projection_matrix, -self.projection_matrix]) + if self.projection_matrix is not None + else None + ) + + def compute_qp_inequalities( + self, configuration: mink.Configuration, dt: float + ) -> mink.Constraint: + """Bound selected joint displacements without constraining scene joints.""" + if self.projection_matrix is None: + return mink.Constraint() + q = configuration.data.qpos[self.qpos_indices] + return mink.Constraint( + G=self._G, + h=self.gain * np.hstack([self.upper - q, q - self.lower]), + ) class ArmJointLimit(mink.Limit): @@ -143,6 +170,11 @@ def __init__( if self.dof_indices.size else None ) + self._G = ( + np.vstack([self.projection_matrix, -self.projection_matrix]) + if self.projection_matrix is not None + else None + ) self._measured_qpos: np.ndarray | None = None def update_measured_state(self, qpos: npt.ArrayLike) -> None: @@ -169,7 +201,7 @@ def compute_qp_inequalities( if dt <= 0.0: raise ValueError("dt must be positive.") - command_q = configuration.q[self.qpos_indices] + command_q = configuration.data.qpos[self.qpos_indices] max_step = dt * self.max_velocity position_lower = self.position_gain * (self.lower - command_q) position_upper = self.position_gain * (self.upper - command_q) @@ -205,7 +237,7 @@ def compute_qp_inequalities( upper_step = np.minimum(upper_step, dt * upper_velocity) return mink.Constraint( - G=np.vstack([self.projection_matrix, -self.projection_matrix]), + G=self._G, h=np.hstack([upper_step, -lower_step]), ) diff --git a/src/openarm_control/qp/bounded_frame_task.py b/src/openarm_control/qp/bounded_frame_task.py index 8b48d723..8a6cf1f9 100644 --- a/src/openarm_control/qp/bounded_frame_task.py +++ b/src/openarm_control/qp/bounded_frame_task.py @@ -143,16 +143,24 @@ def compute_limited_error(self, configuration: mink.Configuration) -> np.ndarray """Return one substep of the total position and orientation budgets.""" return self._clip_error(self.compute_full_error(configuration)) - def _clip_error(self, full_error: np.ndarray) -> np.ndarray: + def _clip_error( + self, + full_error: np.ndarray, + *, + position: bool = True, + orientation: bool = True, + ) -> np.ndarray: """Limit a previously computed frame error.""" error = full_error.copy() for part, total_limit in ( - (slice(0, 3), self.position_error_limit), - (slice(3, 6), self.orientation_error_limit), + (slice(0, 3), self.position_error_limit if position else 0.0), + (slice(3, 6), self.orientation_error_limit if orientation else 0.0), ): limit = total_limit / self._substeps + if limit <= 0.0: + continue norm = float(np.linalg.norm(error[part])) - if limit > 0.0 and norm > limit: + if norm > limit: error[part] *= limit / norm return error @@ -167,39 +175,45 @@ def compute_jacobian(self, configuration: mink.Configuration) -> np.ndarray: def compute_qp_residual( self, configuration: mink.Configuration, - ) -> tuple[np.ndarray, np.ndarray, float] | None: - """Opt out of Mink's fused residual path so bounding is not bypassed. + ) -> tuple[np.ndarray, np.ndarray, float]: + """Return the bounded request for Mink's fused objective assembly.""" + terms = self._bounded_terms(configuration) + if terms is None: + return self.frame_task.compute_qp_residual(configuration) + return self._weighted_residual(*terms) + + def _bounded_terms( + self, configuration: mink.Configuration + ) -> tuple[np.ndarray, np.ndarray] | None: + bound_position = self.position_error_limit > 0.0 and self.limit_activation > 0.0 + bound_orientation = self.orientation_error_limit > 0.0 + if not bound_position and not bound_orientation: + return None - Mink 1.2's solver prefers a task's residual and only falls back to - :meth:`compute_qp_objective` when it is ``None``. The inherited - ``mink.Task`` residual is built from :meth:`compute_error`, which - deliberately reports the *full* error, so accepting it would silently - drop the bounding this class exists to apply. - """ - return None + full_error, jacobian = self.frame_task._error_and_jacobian(configuration) + error = self._clip_error( + full_error, position=bound_position, orientation=bound_orientation + ) + if bound_position: + error[:3] = full_error[:3] + self.limit_activation * ( + error[:3] - full_error[:3] + ) + return error, jacobian def compute_qp_objective( self, configuration: mink.Configuration, ) -> mink.Objective: """Modulate position by its schedule and always bound orientation.""" - bound_position = self.position_error_limit > 0.0 and self.limit_activation > 0.0 - bound_orientation = self.orientation_error_limit > 0.0 - if not bound_position and not bound_orientation: + terms = self._bounded_terms(configuration) + if terms is None: return self.frame_task.compute_qp_objective(configuration) - - full_error = self.compute_full_error(configuration) - limited_error = self._clip_error(full_error) - error = full_error.copy() - if bound_position: - error[:3] += self.limit_activation * (limited_error[:3] - full_error[:3]) - if bound_orientation: - error[3:] = limited_error[3:] + error, jacobian = terms nv = configuration.model.nv if self._identity is None or self._identity.shape != (nv, nv): self._identity = np.eye(nv, dtype=np.float64) return self._assemble_qp( error, - self.compute_jacobian(configuration), + jacobian, self._identity, ) diff --git a/src/openarm_control/qp/nullspace_posture_task.py b/src/openarm_control/qp/nullspace_posture_task.py index 5f6848a6..50ac1086 100644 --- a/src/openarm_control/qp/nullspace_posture_task.py +++ b/src/openarm_control/qp/nullspace_posture_task.py @@ -16,12 +16,15 @@ from __future__ import annotations +import math + import mink import mujoco import numpy as np import numpy.typing as npt from openarm_control.geometry.jacobian import ( + CachedConfiguration, normalized_arm_jacobian, relative_root_is_independent_of_dofs, ) @@ -31,7 +34,8 @@ def smoothstep_activation(value: float, low: float, high: float) -> float: """Map ``value`` to [0, 1] with zero slope at both thresholds.""" if not 0.0 <= low < high: raise ValueError("Expected 0 <= low < high.") - u = float(np.clip((value - low) / (high - low), 0.0, 1.0)) + u = float((value - low) / (high - low)) + u = 0.0 if u < 0.0 else 1.0 if u > 1.0 else u return u * u * (3.0 - 2.0 * u) @@ -45,7 +49,11 @@ def structural_nullspace_direction( raise ValueError(f"Expected a 6x7 Jacobian, got {jacobian.shape}.") _, singular_values, vh = np.linalg.svd(jacobian, full_matrices=True) - direction = vh[-1].copy() + return _align_direction(vh[-1], previous), singular_values + + +def _align_direction(direction: np.ndarray, previous: np.ndarray | None) -> np.ndarray: + direction = direction.copy() if previous is not None: previous = np.asarray(previous, dtype=np.float64) if previous.shape != (7,): @@ -54,7 +62,7 @@ def structural_nullspace_direction( ) if float(direction @ previous) < 0.0: direction = -direction - return direction, singular_values + return direction class NullspacePostureTask(mink.Task): @@ -121,12 +129,52 @@ def __init__( indices, ) self._identity = np.eye(model.nv, dtype=np.float64) + self._configuration_error = np.empty(model.nv, dtype=np.float64) self._previous_direction: np.ndarray | None = None - - def _compute_terms( - self, configuration: mink.Configuration - ) -> tuple[np.ndarray, np.ndarray]: - normalized_jacobian = normalized_arm_jacobian( + self._svd_partner: NullspacePostureTask | None = None + self._svd_batch: np.ndarray | None = None + + def _batch_svd_with(self, other: NullspacePostureTask) -> None: + """Let this task populate both spectra for a fixed pair in one batch.""" + self._svd_partner = other + self._svd_batch = np.empty((2, 6, 7), dtype=np.float64) + + def _svd(self, configuration: mink.Configuration) -> tuple[np.ndarray, np.ndarray]: + cache = ( + configuration._arm_svds + if isinstance(configuration, CachedConfiguration) + else None + ) + cached = cache.get(self) if cache is not None else None + if cached is not None: + return cached + if ( + cache is not None + and self._svd_partner is not None + and self._svd_partner not in cache + ): + assert self._svd_batch is not None + tasks = (self, self._svd_partner) + for index, task in enumerate(tasks): + self._svd_batch[index] = task._jacobian(configuration) + _, singular_values, vh = np.linalg.svd(self._svd_batch, full_matrices=True) + singular_values.setflags(write=False) + vh.setflags(write=False) + for index, task in enumerate(tasks): + cache[task] = (singular_values[index], vh[index, -1]) + return cache[self] + _, singular_values, vh = np.linalg.svd( + self._jacobian(configuration), full_matrices=True + ) + direction = vh[-1] + singular_values.setflags(write=False) + direction.setflags(write=False) + if cache is not None: + cache[self] = (singular_values, direction) + return singular_values, direction + + def _jacobian(self, configuration: mink.Configuration) -> np.ndarray: + return normalized_arm_jacobian( self._frame_task, configuration, self._dof_indices, @@ -134,10 +182,18 @@ def _compute_terms( root_is_independent_of_dofs=self._root_is_independent_of_dofs, ) - direction, singular_values = structural_nullspace_direction( - normalized_jacobian, self._previous_direction - ) - self._previous_direction = direction.copy() + def compute_singularity_ratio(self, configuration: mink.Configuration) -> float: + """Expose the shared spectrum without advancing direction-continuity state.""" + singular_values, _ = self._svd(configuration) + largest = float(singular_values[0]) if singular_values.size else 0.0 + return float(singular_values[-1] / largest) if largest > 0.0 else 0.0 + + def _compute_terms( + self, configuration: mink.Configuration + ) -> tuple[np.ndarray, np.ndarray]: + singular_values, direction = self._svd(configuration) + direction = _align_direction(direction, self._previous_direction) + self._previous_direction = direction largest = float(singular_values[0]) if singular_values.size else 0.0 ratio = float(singular_values[-1] / largest) if largest > 0.0 else 0.0 @@ -145,25 +201,23 @@ def _compute_terms( ratio, self._singularity_low, self._singularity_high ) - configuration_error = np.empty(self._model.nv, dtype=np.float64) + configuration_error = self._configuration_error mujoco.mj_differentiatePos( m=self._model, qvel=configuration_error, dt=1.0, qpos1=self._home_qpos, - qpos2=configuration.q, + qpos2=configuration.data.qpos, ) posture_error = float(direction @ configuration_error[self._dof_indices]) - return_speed = float( - np.clip( - -self._return_rate * posture_error, - -self._max_speed, - self._max_speed, - ) - ) + return_speed = -self._return_rate * posture_error + if return_speed < -self._max_speed: + return_speed = -self._max_speed + elif return_speed > self._max_speed: + return_speed = self._max_speed displacement = return_speed * self._dt - effective_cost = float(np.sqrt(activation) * self._base_cost) + effective_cost = math.sqrt(activation) * self._base_cost self.cost[0] = effective_cost jacobian = np.zeros((1, self._model.nv), dtype=np.float64) jacobian[0, self._dof_indices] = direction @@ -189,14 +243,7 @@ def compute_qp_objective(self, configuration: mink.Configuration) -> mink.Object def compute_qp_residual( self, configuration: mink.Configuration, - ) -> tuple[np.ndarray, np.ndarray, float] | None: - """Fall back to the single-pass objective assembly on Mink 1.2. - - Mink 1.2's solver prefers a task's residual and only falls back to - :meth:`compute_qp_objective` when it is ``None``. The inherited - ``mink.Task`` residual calls :meth:`compute_error` and - :meth:`compute_jacobian` separately, causing each to compute the - normalized Jacobian and SVD. Returning ``None`` preserves this class's - one-SVD :meth:`compute_qp_objective` path. - """ - return None + ) -> tuple[np.ndarray, np.ndarray, float]: + """Compute the fused residual from a single nullspace SVD.""" + error, jacobian = self._compute_terms(configuration) + return self._weighted_residual(error, jacobian) diff --git a/src/openarm_control/qp/singularity_approach_limit.py b/src/openarm_control/qp/singularity_approach_limit.py index 1133f14d..1a7882aa 100644 --- a/src/openarm_control/qp/singularity_approach_limit.py +++ b/src/openarm_control/qp/singularity_approach_limit.py @@ -22,6 +22,7 @@ import numpy.typing as npt from openarm_control.geometry.jacobian import ( + frame_is_independent_of_dofs, normalized_arm_jacobian, relative_root_is_independent_of_dofs, singularity_ratio, @@ -66,6 +67,11 @@ def __init__( self.model = model self.frame_task = frame_task self.dof_indices = indices.copy() + self._jacobian_columns = ( + slice(int(indices[0]), int(indices[-1]) + 1) + if np.all(np.diff(indices) == 1) + else self.dof_indices + ) self.characteristic_length = float(characteristic_length) self.ratio_stop = float(ratio_stop) self.ratio_slow = float(ratio_slow) @@ -78,6 +84,20 @@ def __init__( indices, ) self._scratch = mink.Configuration(model) + native_task = getattr(frame_task, "frame_task", frame_task) + self._frame_id = getattr(model, native_task.frame_type)( + native_task.frame_name + ).id + self._jacobian_func = { + "body": mujoco.mj_jacBody, + "site": mujoco.mj_jacSite, + "geom": mujoco.mj_jacGeom, + }[native_task.frame_type] + self._world_jacobian = np.empty((6, model.nv), dtype=np.float64) + self._tangent = np.zeros(model.nv, dtype=np.float64) + self._perturbed_q = np.empty(model.nq, dtype=np.float64) + self._perturbed_jacobians = np.empty((indices.size, 2, 6, indices.size)) + self._gradient = np.empty(indices.size, dtype=np.float64) self._measured_qpos: np.ndarray | None = None self._G: np.ndarray | None = None self._allowed_rate = self.max_rate @@ -95,12 +115,19 @@ def clear_measured_configuration(self) -> None: """Fall back to the command configuration singularity ratio.""" self._measured_qpos = None - def prepare(self, configuration: mink.Configuration) -> None: - """Linearize rho(q) once for the next outer IK control step.""" - command_ratio, _ = self._ratio(configuration) + def prepare( + self, + configuration: mink.Configuration, + *, + command_ratio: float | None = None, + gradient: np.ndarray | None = None, + ) -> None: + """Linearize rho(q) once, optionally reusing a paired gradient.""" + if command_ratio is None: + command_ratio, _ = self._ratio(configuration) measured_ratio: float | None = None if self._measured_qpos is not None: - self._scratch.update(q=self._measured_qpos) + self._update_scratch(self._measured_qpos) measured_ratio, _ = self._ratio(self._scratch) effective_ratio = ( command_ratio @@ -112,12 +139,44 @@ def prepare(self, configuration: mink.Configuration) -> None: smoothstep = unit_margin * unit_margin * (3.0 - 2.0 * unit_margin) activation = float(smoothstep**self.exponent) self._allowed_rate = self.max_rate * activation - gradient = self._finite_difference_gradient(configuration) + if gradient is None: + gradient = self._finite_difference_gradient(configuration) G = np.zeros((1, self.model.nv), dtype=np.float64) G[0, self.dof_indices] = -gradient self._G = G + def can_share_perturbations(self, other: SingularityApproachLimit) -> bool: + """Check both frames and roots for cross-arm kinematic dependencies.""" + if ( + self.model is not other.model + or self.gradient_epsilon != other.gradient_epsilon + or np.intersect1d(self.dof_indices, other.dof_indices).size + ): + return False + for limit, other_dofs in ( + (self, other.dof_indices), + (other, self.dof_indices), + ): + task = getattr(limit.frame_task, "frame_task", limit.frame_task) + if not ( + frame_is_independent_of_dofs( + self.model, task.frame_name, task.frame_type, other_dofs + ) + and relative_root_is_independent_of_dofs( + limit.frame_task, self.model, other_dofs + ) + ): + return False + return True + + def _paired_gradients( + self, configuration: mink.Configuration, other: SingularityApproachLimit + ) -> tuple[np.ndarray, np.ndarray]: + """Compute a pair already validated by can_share_perturbations.""" + gradient = self._finite_difference_gradient(configuration, paired_limit=other) + return gradient, other._gradient + def compute_qp_inequalities( self, configuration: mink.Configuration, @@ -135,6 +194,26 @@ def compute_qp_inequalities( ) def _ratio(self, configuration: mink.Configuration) -> tuple[float, np.ndarray]: + return singularity_ratio(self._jacobian(configuration)) + + def _jacobian( + self, configuration: mink.Configuration, *, out: np.ndarray | None = None + ) -> np.ndarray: + if self._root_is_independent_of_dofs: + jacobian = self._world_jacobian + self._jacobian_func( + self.model, + configuration.data, + jacobian[:3], + jacobian[3:], + self._frame_id, + ) + # Pure frame rotation preserves singular values, also after length scaling. + if out is None: + out = np.empty((6, self.dof_indices.size), dtype=np.float64) + out[:] = jacobian[:, self._jacobian_columns] + out[:3] /= self.characteristic_length + return out jacobian = normalized_arm_jacobian( self.frame_task, configuration, @@ -142,29 +221,69 @@ def _ratio(self, configuration: mink.Configuration) -> tuple[float, np.ndarray]: self.characteristic_length, root_is_independent_of_dofs=self._root_is_independent_of_dofs, ) - return singularity_ratio(jacobian) + if out is not None: + out[:] = jacobian + return out + return jacobian def _finite_difference_gradient( self, configuration: mink.Configuration, + *, + paired_limit: SingularityApproachLimit | None = None, ) -> np.ndarray: - q0 = configuration.q.copy() - tangent = np.zeros(self.model.nv, dtype=np.float64) - gradient = np.empty(self.dof_indices.size, dtype=np.float64) + q0 = configuration.q + tangent = self._tangent + tangent.fill(0.0) + q = self._perturbed_q + jacobians = self._perturbed_jacobians eps = self.gradient_epsilon for index, dof in enumerate(self.dof_indices): tangent[dof] = 1.0 - q_plus = q0.copy() - mujoco.mj_integratePos(self.model, q_plus, tangent, eps) - self._scratch.update(q=q_plus) - ratio_plus, _ = self._ratio(self._scratch) - - q_minus = q0.copy() - mujoco.mj_integratePos(self.model, q_minus, tangent, -eps) - self._scratch.update(q=q_minus) - ratio_minus, _ = self._ratio(self._scratch) - gradient[index] = (ratio_plus - ratio_minus) / (2.0 * eps) + if paired_limit is not None: + tangent[paired_limit.dof_indices[index]] = 1.0 + q[:] = q0 + mujoco.mj_integratePos(self.model, q, tangent, eps) + self._update_scratch(q) + self._jacobian(self._scratch, out=jacobians[index, 0]) + if paired_limit is not None: + paired_limit._jacobian( + self._scratch, out=paired_limit._perturbed_jacobians[index, 0] + ) + + q[:] = q0 + mujoco.mj_integratePos(self.model, q, tangent, -eps) + self._update_scratch(q) + self._jacobian(self._scratch, out=jacobians[index, 1]) tangent[dof] = 0.0 + if paired_limit is not None: + paired_limit._jacobian( + self._scratch, out=paired_limit._perturbed_jacobians[index, 1] + ) + tangent[paired_limit.dof_indices[index]] = 0.0 + gradient = self._gradient_from_samples() + if paired_limit is not None: + paired_limit._gradient_from_samples() return gradient + + def _gradient_from_samples(self) -> np.ndarray: + # Batch the same positive/negative perturbations without changing the stencil. + singular_values = np.linalg.svd(self._perturbed_jacobians, compute_uv=False) + largest = singular_values[..., 0] + ratios = np.divide( + singular_values[..., -1], + largest, + out=np.zeros_like(largest), + where=largest > 0.0, + ) + np.subtract(ratios[:, 0], ratios[:, 1], out=self._gradient) + self._gradient /= 2.0 * self.gradient_epsilon + return self._gradient + + def _update_scratch(self, q: np.ndarray) -> None: + # Ratio evaluation only needs frame transforms and geometric Jacobians. + self._scratch.data.qpos[:] = q + mujoco.mj_kinematics(self.model, self._scratch.data) + mujoco.mj_comPos(self.model, self._scratch.data) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index e8534525..12b9156b 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -18,7 +18,9 @@ import mink import numpy as np +import pytest +from openarm_control import IKParams, Kinematics from openarm_control.geometry.jacobian import ( normalized_arm_jacobian, relative_root_is_independent_of_dofs, @@ -26,6 +28,43 @@ from _support import make_setup +@pytest.mark.parametrize("operation", ("update", "integrate")) +def test_geometry_cache_is_read_only_and_invalidated(operation: str) -> None: + setup = make_setup("right") + solver = Kinematics(setup, IKParams())._ik + assert solver is not None + configuration = solver._config + task = solver._nullspace_tasks["right"] + frame_name, frame_type = "right_ee_control_point", "site" + dofs = setup.joint_resolver.arm_dof_indices("right") + jacobian = configuration.get_frame_jacobian(frame_name, frame_type) + assert configuration.get_frame_jacobian(frame_name, frame_type) is jacobian + spectrum, direction = task._svd(configuration) + solver._tasks["right"].set_target(mink.SE3.identity()) + assert task._svd(configuration)[0] is spectrum + assert not jacobian.flags.writeable + assert not spectrum.flags.writeable + assert not direction.flags.writeable + + if operation == "update": + q = configuration.q + q[setup.joint_resolver.arm_qpos_indices("right")[3]] += 0.1 + configuration.update(q=q) + else: + velocity = np.zeros(setup.model.nv) + velocity[dofs[3]] = 0.1 + configuration.integrate_inplace(velocity, dt=1.0) + new_jacobian = configuration.get_frame_jacobian(frame_name, frame_type) + new_spectrum, _ = task._svd(configuration) + assert new_jacobian is not jacobian + assert new_spectrum is not spectrum + reference = mink.Configuration(setup.model, q=configuration.q) + np.testing.assert_allclose( + new_jacobian, reference.get_frame_jacobian(frame_name, frame_type) + ) + np.testing.assert_allclose(new_spectrum, task._svd(reference)[0]) + + def test_relative_root_fast_path_matches_full_formula() -> None: setup = make_setup("right", origin_frame="arm_origin") configuration = mink.Configuration( diff --git a/tests/test_ik_params.py b/tests/test_ik_params.py index 99348277..1be501aa 100644 --- a/tests/test_ik_params.py +++ b/tests/test_ik_params.py @@ -71,6 +71,7 @@ def test_unexposed_parameters_remain_regular_ik_fields() -> None: ) assert "--ik-profile" not in options + assert "--diag-reg" not in options assert "--target-linear-speed-slow" not in options assert "--joint-braking-exponent" not in options assert "--singularity-ratio-stop" not in options @@ -127,13 +128,6 @@ def test_control_overrides_and_velocity_yaml(tmp_path: pathlib.Path) -> None: ) -def test_diag_reg_is_not_a_registered_solver_option() -> None: - parser = argparse.ArgumentParser() - register_ik_args(parser) - options = {option for action in parser._actions for option in action.option_strings} - assert "--diag-reg" not in options - - def test_braking_can_be_disabled_without_disabling_velocity_limits() -> None: parser = argparse.ArgumentParser() register_ik_args(parser) diff --git a/tests/test_kinematics.py b/tests/test_kinematics.py index 1efdc271..df7d7f89 100644 --- a/tests/test_kinematics.py +++ b/tests/test_kinematics.py @@ -19,6 +19,7 @@ from unittest import mock import mink +import mujoco import numpy as np import openarm_mujoco.v2 as openarm_mujoco import pytest @@ -29,6 +30,26 @@ from _support import driver_state, make_setup, velocity_mapping +@pytest.mark.parametrize("measured", (False, True)) +def test_solver_batches_svd_and_shares_perturbation_fk(measured: bool) -> None: + setup = make_setup(origin_frame="arm_origin") + params = IKParams(velocity_limits=velocity_mapping(*setup.sides)) + kinematics = Kinematics(setup, params) + for side in setup.sides: + kinematics.set_target(side, setup.read_ee_pose(side)) + if measured: + kinematics.update_measured_state(driver_state(setup)) + with ( + mock.patch.object(np.linalg, "svd", wraps=np.linalg.svd) as svd, + mock.patch.object( + mujoco, "mj_kinematics", wraps=mujoco.mj_kinematics + ) as forward, + ): + assert kinematics.solve() is not None + assert svd.call_count == params.max_iters + 2 + 2 * measured + assert forward.call_count == params.max_iters + 14 + 2 * measured + + def test_driver_qpos_mapping_updates_only_active_arm() -> None: setup = make_setup("right") base_qpos = setup.data.qpos.copy() diff --git a/tests/test_qp_limits.py b/tests/test_qp_limits.py index adf3fbaf..cebf43a9 100644 --- a/tests/test_qp_limits.py +++ b/tests/test_qp_limits.py @@ -18,6 +18,8 @@ import contextlib import io +import xml.etree.ElementTree as ET +from unittest import mock import mink import mujoco @@ -25,6 +27,10 @@ import pytest from openarm_control import ArmSetup, pose_to_se3 +from openarm_control.geometry.jacobian import ( + normalized_arm_jacobian, + singularity_ratio, +) from openarm_control.qp.arm_joint_limit import ( ArmConfigurationLimit, ArmJointLimit, @@ -62,6 +68,7 @@ def test_center_uses_physical_velocity_cap() -> None: configuration.update(q=q) constraint = limit.compute_qp_inequalities(configuration, dt=0.01) assert constraint.h is not None + assert constraint.G is limit.compute_qp_inequalities(configuration, dt=0.01).G assert constraint.h[row] == pytest.approx(limit.max_velocity[row] * 0.01) assert constraint.h[limit.indices.size + row] == pytest.approx( @@ -69,19 +76,32 @@ def test_center_uses_physical_velocity_cap() -> None: ) -def test_configuration_limit_uses_selected_dofs_and_gain() -> None: - setup = make_setup("right") - limit = ArmConfigurationLimit( - setup.model, - setup.joint_resolver.arm_qpos_indices("right"), - gain=0.8, - ) +def test_configuration_limit_selects_scalar_joints_and_applies_gain() -> None: + model = mujoco.MjModel.from_xml_string(""" + + + + + + + + + """) + configuration = mink.Configuration(model) + configuration.update(q=np.array([0.25, np.cos(0.3), np.sin(0.3), 0, 0, -0.5])) + limit = ArmConfigurationLimit(model, [5, 0], gain=0.8) + constraint = limit.compute_qp_inequalities(configuration, dt=0.004) + np.testing.assert_array_equal(limit.indices, [0, 4]) + np.testing.assert_array_equal(limit.qpos_indices, [0, 5]) np.testing.assert_array_equal( - limit.indices, - setup.joint_resolver.arm_dof_indices("right"), + constraint.G, + [[1, 0, 0, 0, 0], [0, 0, 0, 0, 1], [-1, 0, 0, 0, 0], [0, 0, 0, 0, -1]], ) - assert limit.gain == 0.8 + np.testing.assert_allclose(constraint.h, [0.6, 2.0, 1.0, 1.2]) + assert constraint.G is limit.compute_qp_inequalities(configuration, dt=0.004).G + empty = ArmConfigurationLimit(model, [], gain=0.8) + assert empty.compute_qp_inequalities(configuration, dt=0.004).inactive def test_half_braking_distance_allows_quarter_velocity() -> None: @@ -158,16 +178,18 @@ def test_unlimited_selected_joint_prints_warning_and_is_skipped() -> None: assert skipped_qpos not in limit.qpos_indices -def _make_singularity_limit() -> tuple[ - ArmSetup, SingularityApproachLimit, mink.Configuration -]: +def _make_singularity_limit( + root: str | None = None, +) -> tuple[ArmSetup, SingularityApproachLimit, mink.Configuration]: setup = make_setup("right") configuration = mink.Configuration(setup.model, q=setup.data.qpos.copy()) - task = mink.FrameTask( - "right_ee_control_point", - "site", + task_type = mink.FrameTask if root is None else mink.RelativeFrameTask + task = task_type( + frame_name="right_ee_control_point", + frame_type="site", position_cost=10.0, orientation_cost=1.0, + **({"root_name": root, "root_type": "body"} if root is not None else {}), ) limit = SingularityApproachLimit( setup.model, @@ -182,18 +204,81 @@ def _make_singularity_limit() -> tuple[ return setup, limit, configuration +def _scalar_singularity_gradient( + limit: SingularityApproachLimit, configuration: mink.Configuration +) -> np.ndarray: + model = configuration.model + q0 = configuration.q + reference = mink.Configuration(model) + gradient = np.empty(limit.dof_indices.size) + tangent = np.zeros(model.nv) + for index, dof in enumerate(limit.dof_indices): + tangent[dof] = 1.0 + ratios = [] + for sign in (1.0, -1.0): + q = q0.copy() + mujoco.mj_integratePos(model, q, tangent, sign * limit.gradient_epsilon) + reference.update(q=q) + jacobian = normalized_arm_jacobian( + limit.frame_task, + reference, + limit.dof_indices, + limit.characteristic_length, + ) + ratios.append(singularity_ratio(jacobian)[0]) + gradient[index] = (ratios[0] - ratios[1]) / (2.0 * limit.gradient_epsilon) + tangent[dof] = 0.0 + return gradient + + def test_only_approaching_gradient_component_is_bounded() -> None: _, limit, configuration = _make_singularity_limit() + gradient = _scalar_singularity_gradient(limit, configuration) limit.prepare(configuration) constraint = limit.compute_qp_inequalities(configuration, dt=0.004) assert constraint.G is not None - gradient = -constraint.G[0, limit.dof_indices] approach = np.zeros(configuration.model.nv) approach[limit.dof_indices] = -gradient assert float((constraint.G @ approach)[0]) > 0.0 assert float((constraint.G @ (-approach))[0]) < 0.0 +@pytest.mark.parametrize("straight", (False, True)) +@pytest.mark.parametrize("root", (None, "world", "openarm_right_link4")) +def test_batched_gradient_matches_scalar_differences( + straight: bool, root: str | None +) -> None: + setup, limit, configuration = _make_singularity_limit(root) + q0 = configuration.q + if straight: + q0[setup.joint_resolver.arm_qpos_indices("right")] = 0.0 + configuration.update(q=q0) + actual = limit._finite_difference_gradient(configuration).copy() + expected = _scalar_singularity_gradient(limit, configuration) + np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-11) + np.testing.assert_array_equal(configuration.q, q0) + + +def test_scratch_update_preserves_measured_state_constraint() -> None: + setup, limit, configuration = _make_singularity_limit() + q = configuration.q + q[setup.joint_resolver.arm_qpos_indices("right")] += np.linspace(-0.1, 0.1, 7) + configuration.update(q=q) + limit.update_measured_configuration(setup.data.qpos) + + with mock.patch.object(mujoco, "mj_makeConstraint") as make_constraint: + limit.prepare(configuration) + make_constraint.assert_not_called() + actual = limit.compute_qp_inequalities(configuration, dt=0.004) + + with mock.patch.object(limit, "_update_scratch", limit._scratch.update): + limit.prepare(configuration) + reference = limit.compute_qp_inequalities(configuration, dt=0.004) + np.testing.assert_allclose(actual.G, reference.G, rtol=1e-12, atol=1e-12) + np.testing.assert_allclose(actual.h, reference.h, rtol=1e-12, atol=1e-12) + np.testing.assert_array_equal(configuration.q, q) + + def test_target_does_not_change_geometric_ratio() -> None: setup, limit, configuration = _make_singularity_limit() limit.prepare(configuration) @@ -235,3 +320,117 @@ def test_target_does_not_change_geometric_ratio() -> None: np.testing.assert_allclose(shifted.G, initial.G, atol=1e-12) np.testing.assert_allclose(shifted.h, initial.h, atol=1e-12) + + +@pytest.mark.parametrize("reverse", (False, True)) +def test_jacobian_out_preserves_values_and_owned_results(reverse: bool) -> None: + setup, limit, configuration = _make_singularity_limit() + indices = limit.dof_indices.copy() + if reverse: + indices = indices[::-1] + limit = SingularityApproachLimit( + setup.model, + limit.frame_task, + indices, + characteristic_length=0.3, + ratio_stop=0.02, + ratio_slow=0.08, + max_approach_rate=0.25, + ) + owned = limit._jacobian(configuration) + saved = owned.copy() + indices[:] = 0 + output = np.empty((6, 7)) + assert limit._jacobian(configuration, out=output) is output + np.testing.assert_array_equal(output, owned) + assert not np.shares_memory(owned, limit._world_jacobian) + q = configuration.q + q[setup.joint_resolver.arm_qpos_indices("right")] += 0.1 + configuration.update(q=q) + limit._jacobian(configuration, out=output) + np.testing.assert_array_equal(owned, saved) + + +def _make_articulated_pair() -> tuple[ + SingularityApproachLimit, SingularityApproachLimit, mink.Configuration +]: + xml = ET.Element("mujoco") + yaw = ET.SubElement(ET.SubElement(xml, "worldbody"), "body") + ET.SubElement(yaw, "joint", name="yaw", axis="0 0 1") + ET.SubElement(yaw, "geom", type="sphere", size="0.02", mass="1") + base = ET.SubElement(yaw, "body") + ET.SubElement(base, "joint", name="pitch", axis="0 1 0") + ET.SubElement(base, "geom", type="sphere", size="0.02", mass="1") + ET.SubElement(base, "site", name="arm_origin") + for side, offset in (("right", "0 -0.2 0"), ("left", "0 0.2 0")): + body = ET.SubElement(base, "body", pos=offset) + for index in range(7): + body = ET.SubElement( + body, "body", name=f"{side}_link{index}", pos="0 0 0.1" + ) + ET.SubElement( + body, + "joint", + name=f"{side}_joint{index}", + axis=("0 0 1", "0 1 0", "1 0 0")[index % 3], + ) + ET.SubElement(body, "geom", type="sphere", size="0.02", mass="0.1") + ET.SubElement(body, "site", name=f"{side}_tip", pos="0 0 0.1") + model = mujoco.MjModel.from_xml_string(ET.tostring(xml, encoding="unicode")) + limits = [] + for side in ("right", "left"): + task = mink.RelativeFrameTask( + frame_name=f"{side}_tip", + frame_type="site", + root_name="arm_origin", + root_type="site", + position_cost=1.0, + orientation_cost=1.0, + ) + limits.append( + SingularityApproachLimit( + model, + task, + [model.joint(f"{side}_joint{i}").dofadr[0] for i in range(7)], + characteristic_length=0.3, + ratio_stop=0.02, + ratio_slow=0.08, + max_approach_rate=0.25, + ) + ) + return limits[0], limits[1], mink.Configuration(model) + + +@pytest.mark.parametrize("base_angles", ((0.0, 0.0), (0.6, -0.4))) +def test_paired_gradients_preserve_articulated_base( + base_angles: tuple[float, float], +) -> None: + right, left, configuration = _make_articulated_pair() + q = np.linspace(-0.5, 0.6, configuration.model.nq) + q[:2] = base_angles + configuration.update(q=q) + expected = [ + limit._finite_difference_gradient(configuration).copy() + for limit in (right, left) + ] + + assert right.can_share_perturbations(left) + actual = right._paired_gradients(configuration, left) + np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-11) + np.testing.assert_array_equal(configuration.q, q) + + +@pytest.mark.parametrize("change", ("frame", "root", "dof", "epsilon")) +def test_pairing_rejects_cross_dependencies_or_different_steps(change: str) -> None: + right, left, _ = _make_articulated_pair() + if change == "frame": + left.frame_task.frame_name = "right_tip" + elif change == "root": + left.frame_task.root_name = "right_link3" + left.frame_task.root_type = "body" + elif change == "dof": + left.dof_indices[0] = right.dof_indices[0] + else: + left.gradient_epsilon *= 2.0 + assert not right.can_share_perturbations(left) + assert not left.can_share_perturbations(right) diff --git a/tests/test_qp_tasks.py b/tests/test_qp_tasks.py index b07971f4..425e8fe2 100644 --- a/tests/test_qp_tasks.py +++ b/tests/test_qp_tasks.py @@ -121,19 +121,86 @@ def test_orientation_limit_is_independent_of_position_activation() -> None: np.testing.assert_allclose(actual.c, expected.c) -def test_objective_computes_native_error_once() -> None: - configuration, native, task, target = _make_task(orientation_error_limit=0.20) - target[:3] += [0.03, -0.02, 0.01] - task.set_target(pose_to_se3(target)) - - with mock.patch.object( - native, - "compute_error", - wraps=native.compute_error, - ) as compute_error: - task.compute_qp_objective(configuration) - - compute_error.assert_called_once_with(configuration) +def test_clipping_preserves_input_and_disabled_orientation() -> None: + _, _, task, _ = _make_task() + original = np.array([0.1, 0.0, 0.0, 0.2, 0.3, 0.4]) + saved = original.copy() + clipped = task._clip_error(original) + assert not np.shares_memory(clipped, original) + np.testing.assert_array_equal(original, saved) + np.testing.assert_array_equal(clipped[3:], original[3:]) + + +@pytest.mark.parametrize("origin", ("world", "arm_origin")) +@pytest.mark.parametrize( + "activation,orientation_limit", ((0.0, 0.0), (0.5, 0.2), (1.0, 0.2)) +) +def test_fused_residuals_preserve_objective_without_repeated_terms( + origin: str, activation: float, orientation_limit: float +) -> None: + setup = make_setup("right", origin_frame=origin) + solver = Kinematics( + setup, IKParams(nullspace_ratio_low=0.0, nullspace_ratio_high=1e-9) + )._ik + assert solver is not None + configuration = solver._config + q = configuration.q + q[setup.joint_resolver.arm_qpos_indices("right")] += 0.1 + configuration.update(q=q) + frame = solver._tasks["right"] + assert isinstance(frame, BoundedFrameTask) + frame.orientation_error_limit = orientation_limit + frame.set_limit_activation(activation) + target = setup.read_ee_pose("right").astype(np.float64) + target[:3] += [0.04, -0.03, 0.02] + frame.set_target(pose_to_se3(target)) + nullspace = solver._nullspace_tasks["right"] + tasks = [frame, nullspace, solver._kinetic_energy_task] + error = frame.compute_full_error(configuration) + limited = frame._clip_error(error) + error[:3] += activation * (limited[:3] - error[:3]) + if orientation_limit > 0.0: + error[3:] = limited[3:] + expected = [ + frame._assemble_qp( + error, frame.compute_jacobian(configuration), np.eye(configuration.nv) + ), + *[task.compute_qp_objective(configuration) for task in tasks[1:]], + ] + + with ( + mock.patch.object( + frame.frame_task, + "_error_and_jacobian", + wraps=frame.frame_task._error_and_jacobian, + ) as frame_terms, + mock.patch.object( + nullspace, "_compute_terms", wraps=nullspace._compute_terms + ) as terms, + mock.patch.object( + BoundedFrameTask, + "compute_qp_objective", + side_effect=AssertionError("Unexpected fallback"), + ), + mock.patch.object( + NullspacePostureTask, + "compute_qp_objective", + side_effect=AssertionError("Unexpected fallback"), + ), + ): + problem = mink.build_ik(configuration, tasks, 0.004, damping=0.1, limits=[]) + + frame_terms.assert_called_once_with(configuration) + terms.assert_called_once_with(configuration) + np.testing.assert_allclose( + problem.P, + sum(task.H for task in expected) + 0.1 * np.eye(configuration.nv), + rtol=1e-12, + atol=1e-12, + ) + np.testing.assert_allclose( + problem.q, sum(task.c for task in expected), rtol=1e-12, atol=1e-12 + ) @pytest.mark.parametrize("substeps", (1, 5, 10)) @@ -220,9 +287,12 @@ def test_structural_direction_is_null_and_sign_continuous() -> None: def test_smooth_activation_has_flat_clamped_endpoints() -> None: + assert smoothstep_activation(-np.inf, 0.02, 0.05) == 0.0 assert smoothstep_activation(0.02, 0.02, 0.05) == 0.0 assert smoothstep_activation(0.035, 0.02, 0.05) == pytest.approx(0.5) assert smoothstep_activation(0.05, 0.02, 0.05) == 1.0 + assert smoothstep_activation(np.inf, 0.02, 0.05) == 1.0 + assert np.isnan(smoothstep_activation(np.nan, 0.02, 0.05)) def test_task_is_nullspace_only_and_caps_return_speed() -> None: @@ -289,3 +359,39 @@ def test_sync_does_not_move_home_reference() -> None: kinematics.sync(np.linspace(-0.2, 0.2, 16, dtype=np.float32)) for side, task in solver._nullspace_tasks.items(): np.testing.assert_array_equal(task._home_qpos, homes[side]) + + +@pytest.mark.parametrize("partner_first", (False, True)) +def test_batched_spectra_preserve_per_arm_results_and_cache( + partner_first: bool, +) -> None: + setup = make_setup() + solver = Kinematics(setup, IKParams())._ik + configuration = solver._config + tasks = tuple(solver._nullspace_tasks.values()) + expected = [ + structural_nullspace_direction(task._jacobian(configuration)) for task in tasks + ] + order = tuple(reversed(tasks)) if partner_first else tasks + for task in order: + task._svd(configuration) + cached = [task._svd(configuration) for task in tasks] + saved = [(s.copy(), d.copy()) for s, d in cached] + for (spectrum, direction), (expected_direction, expected_spectrum) in zip( + cached, expected + ): + np.testing.assert_allclose(spectrum, expected_spectrum, atol=1e-12) + np.testing.assert_allclose( + np.outer(direction, direction), + np.outer(expected_direction, expected_direction), + atol=1e-12, + ) + assert not spectrum.flags.writeable + assert not direction.flags.writeable + q = configuration.q + q[setup.joint_resolver.arm_qpos_indices("right")] += 0.1 + configuration.update(q=q) + tasks[0]._svd(configuration) + for (spectrum, direction), (old_spectrum, old_direction) in zip(cached, saved): + np.testing.assert_array_equal(spectrum, old_spectrum) + np.testing.assert_array_equal(direction, old_direction)