exoskeleton/code/test/test_bilateral_calibration_v3.py

#!/usr/bin/env python3
"""Bilateral v3 environment, force-audit, and stable-contact contracts."""

from pathlib import Path
import copy
import sys
import unittest

import numpy as np


CODE_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(CODE_ROOT))

from analysis.metrics import MetricError, derive_trial_metrics  # noqa: E402
from experiments.executors import (  # noqa: E402
    _bilateral_data_seed_group,
    _bilateral_network_pair_group_id,
    _bilateral_scenario_and_config,
    execute_bilateral_simulation,
)
from experiments.hashing import stable_hash  # noqa: E402
from experiments.plan import build_trial_plan, load_document  # noqa: E402
from experiments.rng import named_seed_record  # noqa: E402
from simulate_closed_loop import Wall  # noqa: E402


CONFIG_ROOT = CODE_ROOT / "config" / "experiments"


def propagation_trial():
    return {
        "method": {"method_id": "proposed_energy"},
        "trajectory": {
            "trajectory_id": "unit_contact_v3",
            "family": "contact_roundtrip",
            "duration_s": 1.0,
            "contact_probe_fraction": 0.03,
        },
        "factors": {
            "map_policy": "source_stamped",
            "environment_profile": {
                "duration": 6.0,
                "mapping_hz": 40.0,
                "wall_fraction": 0.5,
                "stiffness": 3200.0,
                "damping": 25.0,
                "force_limit": 20.0,
                "transition_depth": 0.001,
                "probe_fraction": 0.0,
                "contact_probe_cycles": 0.0,
            },
            "haptic_profile": {
                "feedback_strength": 0.35,
                "energy_min": 0.0,
                "energy_max": 2.0,
                "energy_initial": 1.0,
            },
            # Direct aliases must override the coupled environment profile.
            "duration_s": 5.0,
            "wall_stiffness": 6400.0,
            "probe_fraction": 0.02,
        },
        "seeds": named_seed_record(11, {"test": "bilateral-v3"}),
    }


def paired_network_trial(profile, *, replicate=0):
    trial = copy.deepcopy(propagation_trial())
    trial["replicate"] = replicate
    trial["factors"].update(
        {
            "bilateral_data_root_seed": 8127,
            "bilateral_pair_network_profiles": True,
            "network_common_random_numbers": True,
            "duration_s": 0.12,
            "network_profile": profile,
        }
    )
    return trial


class BilateralCalibrationV3Test(unittest.TestCase):
    def test_environment_profile_and_direct_factors_propagate(self):
        _, config = _bilateral_scenario_and_config(propagation_trial())
        self.assertEqual(config.duration, 5.0)
        self.assertEqual(config.mapping_hz, 40.0)
        self.assertEqual(config.wall_fraction, 0.5)
        self.assertEqual(config.wall_stiffness, 6400.0)
        self.assertEqual(config.wall_damping, 25.0)
        self.assertEqual(config.wall_force_limit, 20.0)
        self.assertEqual(config.wall_transition_depth, 0.001)
        self.assertEqual(config.contact_probe_fraction, 0.02)
        self.assertEqual(config.contact_probe_cycles, 0.0)

    def test_wall_reports_raw_applied_and_saturation_compatibly(self):
        wall = Wall(
            point=np.zeros(3),
            normal=np.array([1.0, 0.0, 0.0]),
            stiffness=100.0,
            damping=10.0,
            force_limit=5.0,
        )
        contact = wall.contact(
            np.array([0.1, 0.0, 0.0]),
            np.array([2.0, 0.0, 0.0]),
        )
        self.assertAlmostEqual(contact.penetration, 0.1)
        self.assertAlmostEqual(contact.force_raw_N, 30.0)
        self.assertAlmostEqual(contact.force_applied_N, 5.0)
        self.assertTrue(contact.saturation_active)
        np.testing.assert_allclose(
            contact.wrench_applied[:3], [-5.0, 0.0, 0.0]
        )

        legacy_wrench, legacy_penetration = wall.wrench(
            np.array([0.1, 0.0, 0.0]),
            np.array([2.0, 0.0, 0.0]),
        )
        np.testing.assert_allclose(legacy_wrench, contact.wrench_applied)
        self.assertEqual(legacy_penetration, contact.penetration)

    def test_v3_grids_are_separate_and_h3_challenge_is_disabled(self):
        stable = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        challenge_spec = load_document(
            CONFIG_ROOT
            / "bilateral_calibration_v3_energy_challenge.json"
        )
        challenge = build_trial_plan(challenge_spec)
        challenge_metrics = load_document(
            CONFIG_ROOT
            / "metrics_bilateral_v3_energy_challenge.json"
        )

        self.assertEqual(stable["pair_count"], 18)
        self.assertEqual(stable["trial_count"], 18)
        self.assertEqual(challenge["pair_count"], 3)
        self.assertEqual(challenge["trial_count"], 3)
        self.assertFalse(challenge_spec["h3_eligible"])
        self.assertNotIn("h3", challenge_metrics["enabled"])
        self.assertIn("not frozen", stable["specification"]["status"])

    def test_haptic_profiles_and_challenge_share_exogenous_seed_groups(self):
        stable = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        challenge = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_energy_challenge.json"
            )
        )

        stable_k3200_rep0 = [
            trial
            for trial in stable["trials"]
            if trial["replicate"] == 0
            and trial["factors"]["environment_profile"]["profile_id"]
            == "stable_candidate_k3200"
        ]
        self.assertEqual(len(stable_k3200_rep0), 3)
        stable_seeds = {
            _bilateral_scenario_and_config(trial)[1].seed
            for trial in stable_k3200_rep0
        }
        self.assertEqual(len(stable_seeds), 1)

        challenge_rep0 = next(
            trial
            for trial in challenge["trials"]
            if trial["replicate"] == 0
        )
        challenge_seed = _bilateral_scenario_and_config(
            challenge_rep0
        )[1].seed
        self.assertEqual(challenge_seed, next(iter(stable_seeds)))

        stable_k6400_rep0 = next(
            trial
            for trial in stable["trials"]
            if trial["replicate"] == 0
            and trial["factors"]["environment_profile"]["profile_id"]
            == "stiff_candidate_k6400"
        )
        self.assertNotEqual(
            _bilateral_scenario_and_config(stable_k6400_rep0)[1].seed,
            challenge_seed,
        )

    def test_default_v3_data_seed_and_hash_remain_exactly_compatible(self):
        plan = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        trial = next(
            entry
            for entry in plan["trials"]
            if entry["replicate"] == 0
            and entry["factors"]["environment_profile"]["profile_id"]
            == "stable_candidate_k3200"
            and entry["factors"]["haptic_profile"]["profile_id"]
            == "gain_035_wide_tank"
        )
        _, config = _bilateral_scenario_and_config(trial)
        group_id, basis, seed_record = _bilateral_data_seed_group(
            trial, config
        )

        self.assertEqual(trial["trial_id"], "trial-e8342e4017a6d7c5")
        self.assertEqual(config.seed, 1920787158)
        self.assertEqual(
            group_id, "bilateral-data-271d823f09a704f7"
        )
        self.assertEqual(
            stable_hash(basis, prefix="bilateral-data-group"),
            "271d823f09a704f78dde74a3c3bb3c3e"
            "0a5f3b3d5bd288a9c67996f4a4e1ce46",
        )
        self.assertEqual(
            stable_hash(
                seed_record, prefix="bilateral-data-seed-record"
            ),
            "6e5e05b8ac7922c4679464a658a053be"
            "c10da966e43b4d09f2fbaed990da2001",
        )
        self.assertEqual(
            seed_record,
            {
                "simulation": [
                    3977025919,
                    244326538,
                    539806858,
                    605936004,
                ]
            },
        )

    def test_network_profiles_can_share_seed_and_pair_group(self):
        clean = {
            "profile_id": "clean",
            "forward_delay_s": 0.0,
            "return_delay_s": 0.004,
            "forward_jitter_s": 0.0,
            "return_jitter_s": 0.0,
            "forward_packet_loss": 0.0,
            "return_packet_loss": 0.0,
            "forward_timeout_s": 0.05,
            "return_timeout_s": 0.05,
        }
        impaired = {
            "profile_id": "impaired",
            "forward_delay_s": 0.012,
            "return_delay_s": 0.024,
            "forward_jitter_s": 0.003,
            "return_jitter_s": 0.005,
            "forward_packet_loss": 0.1,
            "return_packet_loss": 0.2,
            "forward_timeout_s": 0.08,
            "return_timeout_s": 0.10,
        }
        clean_trial = paired_network_trial(clean)
        impaired_trial = paired_network_trial(impaired)
        _, clean_config = _bilateral_scenario_and_config(clean_trial)
        _, impaired_config = _bilateral_scenario_and_config(impaired_trial)
        clean_group = _bilateral_data_seed_group(
            clean_trial, clean_config
        )
        impaired_group = _bilateral_data_seed_group(
            impaired_trial, impaired_config
        )

        self.assertNotEqual(
            clean_config.feedback_delay_s,
            impaired_config.feedback_delay_s,
        )
        self.assertEqual(clean_config.seed, impaired_config.seed)
        self.assertEqual(clean_group, impaired_group)
        clean_pair_id = _bilateral_network_pair_group_id(
            clean_trial, clean_group[1]
        )
        impaired_pair_id = _bilateral_network_pair_group_id(
            impaired_trial, impaired_group[1]
        )
        self.assertEqual(clean_pair_id, impaired_pair_id)
        self.assertTrue(
            clean_group[1]["network_pairing_strategy"][
                "network_common_random_numbers"
            ]
        )
        self.assertNotIn(
            "forward_delay_s",
            clean_group[1]["mechanical_and_control_inputs"],
        )

        replicate_trial = paired_network_trial(clean, replicate=1)
        _, replicate_config = _bilateral_scenario_and_config(
            replicate_trial
        )
        replicate_group = _bilateral_data_seed_group(
            replicate_trial, replicate_config
        )
        self.assertNotEqual(clean_config.seed, replicate_config.seed)
        self.assertNotEqual(clean_group[0], replicate_group[0])
        self.assertNotEqual(
            clean_pair_id,
            _bilateral_network_pair_group_id(
                replicate_trial, replicate_group[1]
            ),
        )

    def test_network_pairing_requires_common_random_numbers(self):
        trial = paired_network_trial({"profile_id": "invalid_pairing"})
        trial["factors"]["network_common_random_numbers"] = False
        with self.assertRaisesRegex(
            ValueError,
            "bilateral_pair_network_profiles requires "
            "network_common_random_numbers=true",
        ):
            _bilateral_scenario_and_config(trial)

    def test_nominal_transport_matches_within_trial_shadows(self):
        nominal = {
            "profile_id": "nominal",
            "forward_delay_s": 0.0,
            "return_delay_s": 0.0,
            "forward_jitter_s": 0.0,
            "return_jitter_s": 0.0,
            "forward_packet_loss": 0.0,
            "return_packet_loss": 0.0,
            "forward_timeout_s": 0.04,
            "return_timeout_s": 0.04,
        }
        samples = execute_bilateral_simulation(
            paired_network_trial(nominal)
        ).samples
        np.testing.assert_array_equal(
            samples["q_slave_ref"],
            samples["q_slave_zero_delay_ref"],
        )
        np.testing.assert_array_equal(
            samples["tau_master_mapped"],
            samples["tau_master_zero_return_delay"],
        )
        np.testing.assert_array_equal(
            samples["slave_reference_lag_error"],
            np.zeros(samples["time"].shape[0]),
        )
        np.testing.assert_array_equal(
            samples["return_feedback_lag_error"],
            np.zeros(samples["time"].shape[0]),
        )

    def test_short_network_run_exports_pairing_and_shadow_diagnostics(self):
        profile = {
            "profile_id": "short_delay",
            "forward_delay_s": 0.006,
            "return_delay_s": 0.008,
            "forward_jitter_s": 0.001,
            "return_jitter_s": 0.001,
            "forward_packet_loss": 0.0,
            "return_packet_loss": 0.0,
            "forward_timeout_s": 0.04,
            "return_timeout_s": 0.04,
        }
        payload = execute_bilateral_simulation(
            paired_network_trial(profile)
        )
        samples = payload.samples
        sample_count = samples["time"].shape[0]
        self.assertEqual(sample_count, 60)

        for field in (
            "forward_packet_active",
            "forward_packet_fresh",
            "return_packet_fresh",
            "forward_packet_age",
            "forward_packet_seq",
            "return_packet_seq",
            "slave_zero_delay_tracking_error",
            "slave_reference_lag_error",
            "return_feedback_lag_error",
        ):
            self.assertEqual(samples[field].shape, (sample_count,))
        for field in (
            "q_slave_zero_delay_ref",
            "tau_master_zero_return_delay",
        ):
            self.assertEqual(samples[field].shape, (sample_count, 7))
            self.assertTrue(np.all(np.isfinite(samples[field])))

        for direction in ("forward", "return"):
            active = samples[f"{direction}_packet_active"].astype(bool)
            fresh = samples[f"{direction}_packet_fresh"].astype(bool)
            age = samples[f"{direction}_packet_age"]
            seq = samples[f"{direction}_packet_seq"]
            self.assertTrue(np.any(active))
            self.assertTrue(np.any(~active))
            self.assertTrue(np.all(~fresh | active))
            self.assertTrue(np.all(np.isnan(age[~active])))
            self.assertTrue(np.all(np.isfinite(age[active])))
            self.assertTrue(np.all(age[active] >= 0.0))
            self.assertTrue(np.all(seq[~active] == -1))
            self.assertTrue(np.all(seq[active] >= 0))

        for field in (
            "slave_zero_delay_tracking_error",
            "slave_reference_lag_error",
            "return_feedback_lag_error",
        ):
            self.assertTrue(np.all(np.isfinite(samples[field])))
            self.assertTrue(np.all(samples[field] >= 0.0))
        np.testing.assert_allclose(
            samples["return_feedback_lag_error"],
            np.linalg.norm(
                samples["tau_master_zero_return_delay"]
                - samples["tau_master_mapped"],
                axis=1,
            ),
        )

        configured = {
            "configured_forward_delay_s": 0.006,
            "configured_return_delay_s": 0.008,
            "configured_forward_jitter_s": 0.001,
            "configured_return_jitter_s": 0.001,
            "configured_forward_packet_loss": 0.0,
            "configured_return_packet_loss": 0.0,
            "configured_forward_timeout_s": 0.04,
            "configured_return_timeout_s": 0.04,
            "configured_mapping_hz": 40.0,
        }
        for field, expected in configured.items():
            np.testing.assert_array_equal(
                samples[field], np.full(sample_count, expected)
            )
        np.testing.assert_array_equal(
            samples["contact_expected"],
            np.ones(sample_count, dtype=np.int8),
        )
        np.testing.assert_array_equal(
            samples["configured_network_common_random_numbers"],
            np.ones(sample_count, dtype=np.int8),
        )

        self.assertTrue(payload.metadata["network_pairing_enabled"])
        self.assertIn("network_pair_group_id", payload.metadata)
        self.assertEqual(
            payload.metadata["bilateral_data_seed_record_hash"],
            stable_hash(
                payload.metadata["bilateral_data_seed_record"],
                prefix="bilateral-data-seed-record",
            ),
        )

    def test_stable_six_second_smoke_has_contact_without_force_limiting(self):
        plan = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_stable_contact.json"
            )
        )
        trial = next(
            entry
            for entry in plan["trials"]
            if entry["factors"]["environment_profile"]["profile_id"]
            == "stable_candidate_k3200"
            and entry["factors"]["haptic_profile"]["profile_id"]
            == "gain_035_wide_tank"
        )
        payload = execute_bilateral_simulation(trial)
        applied = payload.samples["wall_force_applied_N"]

        self.assertEqual(applied.shape[0], 3000)
        self.assertTrue(
            np.all(payload.samples["wall_force_saturation_active"] == 0)
        )
        self.assertGreaterEqual(float(np.max(applied)), 0.5)
        self.assertLessEqual(float(np.max(applied)), 5.0)
        np.testing.assert_allclose(
            payload.samples["wall_force_raw_N"],
            payload.samples["wall_force_applied_N"],
        )

        metric_config = load_document(
            CONFIG_ROOT
            / "metrics_bilateral_v3_stable_contact.json"
        )
        metrics = derive_trial_metrics(
            payload.samples,
            metric_config,
        )
        self.assertTrue(metrics["bilateral_stable_contact_gate_pass"])
        self.assertEqual(metrics["bilateral_force_limit_hit_fraction"], 0.0)
        self.assertEqual(metrics["bilateral_limit_active_fraction"], 0.0)
        self.assertGreaterEqual(
            metrics["bilateral_contact_force_rms_N"], 0.1
        )
        incomplete = dict(payload.samples)
        incomplete.pop("haptic_rate_limit_active")
        with self.assertRaisesRegex(
            MetricError, "missing required bilateral audit fields"
        ):
            derive_trial_metrics(incomplete, metric_config)

    def test_energy_challenge_is_active_audited_and_unsaturated(self):
        plan = build_trial_plan(
            load_document(
                CONFIG_ROOT
                / "bilateral_calibration_v3_energy_challenge.json"
            )
        )
        payload = execute_bilateral_simulation(plan["trials"][0])
        metrics = derive_trial_metrics(
            payload.samples,
            load_document(
                CONFIG_ROOT
                / "metrics_bilateral_v3_energy_challenge.json"
            ),
        )

        self.assertTrue(metrics["h4_energy_audit_pass"])
        self.assertTrue(metrics["energy_challenge_gate_pass"])
        self.assertGreater(metrics["h4_shadow_floor_deficit_J"], 0.0)
        self.assertGreater(
            metrics["bilateral_energy_probe_raw_work_J"], 0.0
        )
        self.assertGreaterEqual(
            metrics["bilateral_projection_intervention_fraction"], 0.02
        )
        self.assertLessEqual(
            metrics["bilateral_projection_intervention_fraction"], 0.30
        )
        self.assertEqual(metrics["bilateral_force_limit_hit_fraction"], 0.0)
        self.assertEqual(metrics["bilateral_limit_active_fraction"], 0.0)
        self.assertTrue(metrics["bilateral_stable_contact_gate_pass"])
        self.assertTrue(np.all(payload.samples["h3_eligible"] == 0))
        self.assertFalse(payload.metadata["h3_eligible"])
        self.assertEqual(
            payload.metadata["energy_probe"]["window"],
            "hann_squared_sine",
        )
        with self.assertRaisesRegex(MetricError, "H3 is ineligible"):
            derive_trial_metrics(
                payload.samples,
                {"enabled": ["h3"], "h3": {}},
            )


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