from __future__ import annotations from typing import Optional import numpy as np import pinocchio as pin from .wrench_solver import WrenchSolveResult, WrenchSolver def checked_frame_id(model: pin.Model, frame_name: str) -> int: """Resolve a Pinocchio frame name and reject its not-found sentinel.""" if not isinstance(frame_name, str) or not frame_name: raise ValueError("frame_name must be a non-empty string") frame_id = int(model.getFrameId(frame_name)) if frame_id < 0 or frame_id >= model.nframes: available = ", ".join(frame.name for frame in model.frames) raise ValueError( f"Frame not found: {frame_name!r}. " f"Expected an id in [0, {model.nframes}), got {frame_id}. " f"Available frames: {available}" ) return frame_id def checked_joint_id(model: pin.Model, joint_name: str) -> int: """Resolve a movable joint name and reject universe/not-found sentinels.""" if not isinstance(joint_name, str) or not joint_name: raise ValueError("joint_name must be a non-empty string") joint_id = int(model.getJointId(joint_name)) if joint_id <= 0 or joint_id >= model.njoints: available = ", ".join(model.names[1:]) raise ValueError( f"Movable joint not found: {joint_name!r}. " f"Expected an id in [1, {model.njoints}), got {joint_id}. " f"Available movable joints: {available}" ) return joint_id def _as_vector(value, size: int, name: str) -> np.ndarray: vector = np.asarray(value, dtype=float).reshape(-1) if vector.shape != (size,): raise ValueError(f"{name} must have shape ({size},), got {vector.shape}") if not np.all(np.isfinite(vector)): raise ValueError(f"{name} must contain only finite values") return vector def _as_sample_matrix(value, width: int, name: str) -> np.ndarray: samples = np.asarray(value, dtype=float) if samples.ndim == 1: samples = samples.reshape(1, -1) if samples.ndim != 2 or samples.shape[1] != width: raise ValueError( f"{name} must have shape (n_samples, {width}), got {samples.shape}" ) if samples.shape[0] == 0: raise ValueError(f"{name} must contain at least one sample") if not np.all(np.isfinite(samples)): raise ValueError(f"{name} must contain only finite values") return samples class InteractionEstimator: """ Estimate an external wrench at the end-effector point. Conventions are explicit throughout this class: * twists are ``[linear_velocity; angular_velocity]``; * wrenches are ``[force; moment]``; * both are expressed along the chest-frame axes, but at the EE point; * ``tau_int = tau_meas - tau_model - tau_bias``. ``estimate`` keeps the original three-value return contract. Its first result is the bias-corrected residual. The raw and corrected residuals are also available through ``last_tau_residual_raw`` and ``last_tau_residual_corrected``. """ def __init__( self, model: pin.Model, chest_frame_name: str, ee_frame_name: str, lambda_damp: float = 1e-3, wrench_solver: Optional[WrenchSolver] = None, ): damping = float(lambda_damp) if not np.isfinite(damping) or damping <= 0.0: raise ValueError("lambda_damp must be a finite positive scalar") self.model = model self.data = model.createData() self.lambda_damp = damping self.wrench_solver = wrench_solver self.fid_C = checked_frame_id(model, chest_frame_name) self.fid_EE = checked_frame_id(model, ee_frame_name) self._tau_bias = np.zeros(model.nv, dtype=float) self._last_tau_residual_raw: Optional[np.ndarray] = None self._last_tau_residual_corrected: Optional[np.ndarray] = None self._last_wrench_solve: Optional[WrenchSolveResult] = None @staticmethod def _rotation6(rotation: np.ndarray) -> np.ndarray: """Apply one 3-D rotation to both linear and angular blocks.""" rotation = np.asarray(rotation, dtype=float) if rotation.shape != (3, 3): raise ValueError(f"rotation must have shape (3, 3), got {rotation.shape}") rotation6 = np.zeros((6, 6), dtype=float) rotation6[:3, :3] = rotation rotation6[3:, 3:] = rotation return rotation6 @property def tau_bias(self) -> np.ndarray: """Current no-contact joint-torque bias (copy).""" return self._tau_bias.copy() @property def last_tau_residual_raw(self) -> Optional[np.ndarray]: """Latest ``tau_meas - tau_model`` sample, before bias removal.""" if self._last_tau_residual_raw is None: return None return self._last_tau_residual_raw.copy() @property def last_tau_residual_corrected(self) -> Optional[np.ndarray]: """Latest raw residual minus the calibrated bias.""" if self._last_tau_residual_corrected is None: return None return self._last_tau_residual_corrected.copy() @property def last_wrench_solve(self) -> Optional[WrenchSolveResult]: """Latest formal scaled-solver result, if one was configured.""" return self._last_wrench_solve def calibrate_bias(self, tau_residual_raw_samples) -> np.ndarray: """ Calibrate a constant joint-torque bias from offline no-contact samples. Parameters ---------- tau_residual_raw_samples: One sample with shape ``(nv,)`` or a batch with shape ``(n_samples, nv)``. Every row must already be the raw residual ``tau_meas - tau_model`` collected under a no-contact condition. Returns ------- numpy.ndarray A copy of the calibrated mean bias. """ samples = _as_sample_matrix( tau_residual_raw_samples, self.model.nv, "tau_residual_raw_samples" ) self._tau_bias = np.mean(samples, axis=0) return self.tau_bias def calibrate_bias_from_measurements( self, q_samples, qd_samples, qdd_samples, tau_meas_samples, tau_ff_fric_samples=None, ) -> np.ndarray: """ Compute and calibrate bias from an offline no-contact measurement batch. Each argument is a row-major sample matrix. ``tau_ff_fric_samples`` is optional and, when supplied, must have the same ``(n_samples, nv)`` shape as the velocity/torque batches. """ q_batch = _as_sample_matrix(q_samples, self.model.nq, "q_samples") qd_batch = _as_sample_matrix(qd_samples, self.model.nv, "qd_samples") qdd_batch = _as_sample_matrix(qdd_samples, self.model.nv, "qdd_samples") tau_batch = _as_sample_matrix( tau_meas_samples, self.model.nv, "tau_meas_samples" ) sample_count = q_batch.shape[0] batches = (qd_batch, qdd_batch, tau_batch) if any(batch.shape[0] != sample_count for batch in batches): raise ValueError("all calibration batches must have the same sample count") friction_batch = None if tau_ff_fric_samples is not None: friction_batch = _as_sample_matrix( tau_ff_fric_samples, self.model.nv, "tau_ff_fric_samples" ) if friction_batch.shape[0] != sample_count: raise ValueError( "tau_ff_fric_samples must match the calibration sample count" ) residuals = np.empty((sample_count, self.model.nv), dtype=float) for sample_index in range(sample_count): friction = ( None if friction_batch is None else friction_batch[sample_index] ) tau_model = self._tau_model( q_batch[sample_index], qd_batch[sample_index], qdd_batch[sample_index], friction, ) residuals[sample_index] = tau_batch[sample_index] - tau_model return self.calibrate_bias(residuals) def clear_bias(self) -> None: """Clear the calibrated joint-torque bias.""" self._tau_bias.fill(0.0) def _chest_jacobian(self, q, qd): """ Return the EE-point Jacobian expressed along chest-frame axes. ``LOCAL_WORLD_ALIGNED`` is essential here: unlike ``WORLD``, its translational rows are the linear velocity of the EE origin. Rotating the two 3-D blocks changes only their coordinate axes; no translational adjoint term is used, so the wrench/twist reference point stays at EE. """ q = _as_vector(q, self.model.nq, "q") qd = _as_vector(qd, self.model.nv, "qd") pin.forwardKinematics(self.model, self.data, q, qd) pin.updateFramePlacements(self.model, self.data) jacobian_lwa = pin.computeFrameJacobian( self.model, self.data, q, self.fid_EE, pin.ReferenceFrame.LOCAL_WORLD_ALIGNED, ) rotation_world_from_chest = self.data.oMf[self.fid_C].rotation rotation_chest_from_world = rotation_world_from_chest.T return self._rotation6(rotation_chest_from_world) @ jacobian_lwa def _tau_model(self, q, qd, qdd, tau_ff_fric=None): """Return ``M qdd + C qd + g + tau_ff_fric``.""" q = _as_vector(q, self.model.nq, "q") qd = _as_vector(qd, self.model.nv, "qd") qdd = _as_vector(qdd, self.model.nv, "qdd") mass_matrix = pin.crba(self.model, self.data, q) mass_matrix = ( mass_matrix + mass_matrix.T - np.diag(mass_matrix.diagonal()) ) nonlinear = pin.nonLinearEffects(self.model, self.data, q, qd) tau_model = mass_matrix @ qdd + nonlinear if tau_ff_fric is not None: tau_model = tau_model + _as_vector( tau_ff_fric, self.model.nv, "tau_ff_fric" ) return tau_model def estimate(self, q, qd, qdd, tau_meas, tau_ff_fric=None): """ Estimate the bias-corrected joint residual and chest-axis EE wrench. A fixed damped least-squares solve is used: ``F = (J J.T + lambda**2 I)^-1 J tau_int``. """ tau_meas = _as_vector(tau_meas, self.model.nv, "tau_meas") tau_model = self._tau_model(q, qd, qdd, tau_ff_fric) tau_residual_raw = tau_meas - tau_model tau_residual_corrected = tau_residual_raw - self._tau_bias self._last_tau_residual_raw = tau_residual_raw.copy() self._last_tau_residual_corrected = tau_residual_corrected.copy() chest_jacobian = self._chest_jacobian(q, qd) if self.wrench_solver is None: normal_matrix = chest_jacobian @ chest_jacobian.T normal_matrix = normal_matrix + ( self.lambda_damp**2 ) * np.eye(6, dtype=float) wrench_chest = np.linalg.solve( normal_matrix, chest_jacobian @ tau_residual_corrected ) self._last_wrench_solve = None else: solve = self.wrench_solver.solve( chest_jacobian, tau_residual_corrected ) wrench_chest = solve.wrench.copy() self._last_wrench_solve = solve return tau_residual_corrected, wrench_chest, chest_jacobian