exoskeleton/code/test/test_interaction_estimator_contract.py

from __future__ import annotations

import sys
import unittest
from pathlib import Path

import numpy as np
import pinocchio as pin


CODE_DIR = Path(__file__).resolve().parents[1]
if str(CODE_DIR) not in sys.path:
    sys.path.insert(0, str(CODE_DIR))

from core.interaction_estimater import (  # noqa: E402
    InteractionEstimator,
    checked_frame_id,
    checked_joint_id,
)


MASTER_URDF = CODE_DIR / "config" / "master_7dof.urdf"
GENERAL_Q = np.array([0.3, -0.5, 0.4, 0.8, -0.2, 0.6, -0.4])


class InteractionEstimatorContractTest(unittest.TestCase):
    def setUp(self) -> None:
        self.model = pin.buildModelFromUrdf(str(MASTER_URDF))

    def test_pinocchio_not_found_sentinels_are_rejected(self) -> None:
        model = self.model
        self.assertLess(checked_frame_id(model, "master_ee"), model.nframes)
        self.assertLess(
            checked_joint_id(model, "master_elbow_flex_joint"), model.njoints
        )

        with self.assertRaisesRegex(ValueError, "Frame not found"):
            InteractionEstimator(model, "missing_chest", "master_ee")
        with self.assertRaisesRegex(ValueError, "Frame not found"):
            InteractionEstimator(model, "master_base", "missing_ee")
        with self.assertRaisesRegex(ValueError, "Movable joint not found"):
            checked_joint_id(model, "missing_joint")

    def test_virtual_work_is_preserved_at_the_ee_point(self) -> None:
        model = self.model
        estimator = InteractionEstimator(
            model, "master_forearm", "master_ee", lambda_damp=1e-4
        )
        joint_velocity = np.array([0.2, -0.1, 0.3, 0.4, -0.2, 0.1, 0.5])
        wrench_chest = np.array([8.0, -3.0, 5.0, 0.7, -0.4, 0.2])

        jacobian_chest = estimator._chest_jacobian(GENERAL_Q, joint_velocity)
        tau_external = jacobian_chest.T @ wrench_chest
        twist_chest = jacobian_chest @ joint_velocity

        np.testing.assert_allclose(
            tau_external @ joint_velocity,
            wrench_chest @ twist_chest,
            rtol=1e-13,
            atol=1e-13,
        )

    def test_chest_axis_rotation_matches_lwa_jacobian(self) -> None:
        model = self.model
        estimator = InteractionEstimator(
            model, "master_forearm", "master_ee", lambda_damp=1e-4
        )
        joint_velocity = np.array([0.1, 0.2, -0.3, 0.05, 0.4, -0.2, 0.1])

        jacobian_chest = estimator._chest_jacobian(GENERAL_Q, joint_velocity)
        jacobian_lwa = pin.computeFrameJacobian(
            model,
            estimator.data,
            GENERAL_Q,
            estimator.fid_EE,
            pin.ReferenceFrame.LOCAL_WORLD_ALIGNED,
        )
        rotation_world_from_chest = estimator.data.oMf[estimator.fid_C].rotation
        rotation_chest_from_world = rotation_world_from_chest.T
        rotation6 = estimator._rotation6(rotation_chest_from_world)

        np.testing.assert_allclose(
            jacobian_chest, rotation6 @ jacobian_lwa, rtol=1e-13, atol=1e-13
        )

        wrench_chest = np.array([4.0, -2.0, 6.0, 0.3, 0.5, -0.1])
        wrench_world = rotation6.T @ wrench_chest
        np.testing.assert_allclose(
            jacobian_chest.T @ wrench_chest,
            jacobian_lwa.T @ wrench_world,
            rtol=1e-13,
            atol=1e-13,
        )

    def test_bias_corrected_ideal_wrench_is_recovered(self) -> None:
        model = self.model
        estimator = InteractionEstimator(
            model, "master_base", "master_ee", lambda_damp=1e-7
        )
        qd = np.array([0.08, -0.04, 0.03, 0.02, -0.05, 0.06, -0.01])
        qdd = np.array([0.2, -0.1, 0.05, 0.08, -0.04, 0.03, -0.02])
        wrench_true = np.array([9.0, -4.0, 6.0, 0.8, -0.3, 0.5])
        bias_true = np.array([0.12, -0.08, 0.05, 0.03, -0.06, 0.04, -0.02])

        jacobian = estimator._chest_jacobian(GENERAL_Q, qd)
        self.assertEqual(np.linalg.matrix_rank(jacobian, tol=1e-9), 6)
        tau_contact = jacobian.T @ wrench_true
        tau_model = estimator._tau_model(GENERAL_Q, qd, qdd)

        calibration_delta = np.linspace(-0.01, 0.01, model.nv)
        calibrated = estimator.calibrate_bias(
            np.vstack(
                [
                    bias_true + calibration_delta,
                    bias_true - calibration_delta,
                    bias_true,
                ]
            )
        )
        np.testing.assert_allclose(calibrated, bias_true, atol=1e-15)

        tau_corrected, wrench_estimated, jacobian_estimated = estimator.estimate(
            GENERAL_Q, qd, qdd, tau_model + bias_true + tau_contact
        )

        np.testing.assert_allclose(jacobian_estimated, jacobian, atol=1e-13)
        np.testing.assert_allclose(
            estimator.last_tau_residual_raw, bias_true + tau_contact, atol=1e-12
        )
        np.testing.assert_allclose(
            estimator.last_tau_residual_corrected, tau_contact, atol=1e-12
        )
        np.testing.assert_allclose(tau_corrected, tau_contact, atol=1e-12)
        np.testing.assert_allclose(
            wrench_estimated, wrench_true, rtol=2e-11, atol=2e-11
        )

        estimator.clear_bias()
        np.testing.assert_array_equal(estimator.tau_bias, np.zeros(model.nv))

    def test_offline_measurement_batch_calibrates_bias(self) -> None:
        model = self.model
        estimator = InteractionEstimator(
            model, "master_base", "master_ee", lambda_damp=1e-4
        )
        bias_true = np.linspace(-0.09, 0.12, model.nv)
        q_batch = np.vstack([GENERAL_Q, GENERAL_Q + 0.03, GENERAL_Q - 0.02])
        qd_batch = np.vstack(
            [
                np.zeros(model.nv),
                np.linspace(-0.02, 0.03, model.nv),
                np.linspace(0.01, -0.04, model.nv),
            ]
        )
        qdd_batch = np.vstack(
            [
                np.zeros(model.nv),
                np.linspace(0.04, -0.02, model.nv),
                np.linspace(-0.03, 0.05, model.nv),
            ]
        )
        tau_meas_batch = np.vstack(
            [
                estimator._tau_model(q, qd, qdd) + bias_true
                for q, qd, qdd in zip(q_batch, qd_batch, qdd_batch)
            ]
        )

        calibrated = estimator.calibrate_bias_from_measurements(
            q_batch, qd_batch, qdd_batch, tau_meas_batch
        )
        np.testing.assert_allclose(calibrated, bias_true, atol=1e-12)

    def test_fixed_dls_remains_finite_at_a_singularity(self) -> None:
        model = self.model
        estimator = InteractionEstimator(
            model, "master_base", "master_ee", lambda_damp=1e-3
        )
        q = np.zeros(model.nq)
        qd = np.zeros(model.nv)
        qdd = np.zeros(model.nv)
        jacobian = estimator._chest_jacobian(q, qd)
        self.assertLess(np.linalg.matrix_rank(jacobian, tol=1e-9), 6)

        tau_model = estimator._tau_model(q, qd, qdd)
        residual = np.linspace(-0.5, 0.7, model.nv)
        tau_corrected, wrench_estimated, jacobian_estimated = estimator.estimate(
            q, qd, qdd, tau_model + residual
        )

        np.testing.assert_allclose(tau_corrected, residual, atol=1e-12)
        np.testing.assert_allclose(jacobian_estimated, jacobian, atol=1e-13)
        self.assertTrue(np.all(np.isfinite(wrench_estimated)))


if __name__ == "__main__":
    unittest.main()