exoskeleton/code/experiments/executors.py

"""Executable pre-prototype G0c studies.

Each function accepts one immutable trial record produced by
``experiments.plan`` and returns an atomic ``TrialPayload``.  The methods in a
pair use the same recorded trajectory/model/sensor/network seeds.
"""

from __future__ import annotations

from dataclasses import replace
from enum import Enum
from typing import Any, Mapping

import numpy as np
import pinocchio as pin

from core.estimation_signals import (
    JointFrictionCalibration,
    ResidualAblation,
    WrenchEstimatorCalibration,
    CalibratedResidualWrenchEstimator,
)
from core.model_contract import (
    MASTER_JOINT_NAMES,
    SLAVE_FRAMES,
    SLAVE_JOINT_NAMES,
    finite_joint_limits,
    load_models,
    require_frame,
)
from core.retargeting_baselines import (
    RetargetingFailure,
    build_canonical_sew_target_baselines,
)
from core.wrench_solver import ScaledDLSSolver, UndampedSVDSolver
from experiments.io import TrialPayload
from experiments.rng import generator_from_record
from simulate_closed_loop import (
    SCENARIOS,
    SimulationConfig,
    build_mapper,
    make_wall,
    simulate_scenario,
)


def _method_id(trial: Mapping[str, Any]) -> str:
    method = trial.get("method")
    if not isinstance(method, Mapping) or not isinstance(
        method.get("method_id"), str
    ):
        raise ValueError("trial has no method.method_id")
    return method["method_id"]


def _trajectory_spec(trial: Mapping[str, Any]) -> Mapping[str, Any]:
    trajectory = trial.get("trajectory")
    if not isinstance(trajectory, Mapping):
        raise ValueError("trial has no trajectory mapping")
    return trajectory


def _factor(trial: Mapping[str, Any], name: str, default: Any) -> Any:
    factors = trial.get("factors", {})
    if not isinstance(factors, Mapping):
        raise ValueError("trial factors must be a mapping")
    return factors.get(name, default)


def _enum_code(member: Enum) -> int:
    return list(type(member)).index(member)


def _master_trajectory(
    trial: Mapping[str, Any],
    *,
    lower: np.ndarray,
    upper: np.ndarray,
) -> np.ndarray:
    """Generate a continuous, bounded master trajectory from a frozen spec."""
    specification = _trajectory_spec(trial)
    sample_count = int(specification.get("sample_count", 81))
    if sample_count < 3:
        raise ValueError("H1 trajectory sample_count must be at least three")
    family = str(specification.get("family", "nominal"))
    center = np.asarray(
        specification.get(
            "center",
            [0.534, 0.314, -0.10, 2.14, 0.38, 0.38, -0.72],
        ),
        dtype=float,
    )
    delta = np.asarray(
        specification.get(
            "delta",
            [0.08, -0.06, 0.05, -0.12, 0.04, 0.05, -0.04],
        ),
        dtype=float,
    )
    if center.shape != (7,) or delta.shape != (7,):
        raise ValueError("H1 center and delta must have seven entries")
    if family == "joint_limit":
        center = center.copy()
        center[0] = upper[0] - 0.03
        delta = np.zeros(7)
        delta[0] = -0.22
    elif family == "low_manipulability":
        center = center.copy()
        center[3] = 0.08
        delta = np.array([0.04, 0.03, -0.04, 0.05, 0.02, -0.02, 0.02])
    elif family == "reach_boundary":
        delta = 1.75 * delta
    elif family == "sew_degeneracy":
        center = np.array([0.0, 0.0, 0.0, 0.12, 0.0, 0.0, 0.0])
        delta = np.array([0.0, 0.18, 0.0, 0.08, 0.0, -0.08, 0.0])

    phase = np.linspace(0.0, 1.0, sample_count)
    # One cosine excursion starts and ends at the same configuration with zero
    # endpoint velocity, making discontinuities attributable to the mapper.
    excursion = 0.5 - 0.5 * np.cos(2.0 * np.pi * phase)
    trajectory = center[None, :] + excursion[:, None] * delta[None, :]
    margin = 1e-4
    if np.any(trajectory < lower + margin) or np.any(trajectory > upper - margin):
        raise ValueError(
            f"trajectory {specification.get('trajectory_id')} exceeds master limits"
        )
    return trajectory


def _slave_swivel(
    model: pin.Model,
    data: pin.Data,
    q_slave: np.ndarray,
    shoulder_id: int,
    elbow_id: int,
    wrist_id: int,
    previous: float,
) -> tuple[float, bool]:
    """Evaluate a continuous diagnostic arm-plane angle from slave geometry."""
    pin.forwardKinematics(model, data, q_slave)
    pin.updateFramePlacements(model, data)
    shoulder = data.oMf[shoulder_id].translation
    elbow = data.oMf[elbow_id].translation
    wrist = data.oMf[wrist_id].translation
    axis = wrist - shoulder
    axis_norm = float(np.linalg.norm(axis))
    if axis_norm <= 1e-9:
        return previous, True
    axis /= axis_norm
    radial = elbow - shoulder
    radial -= float(radial @ axis) * axis
    radial_norm = float(np.linalg.norm(radial))
    if radial_norm <= 1e-9:
        return previous, True
    radial /= radial_norm
    reference = np.array([0.0, 0.0, 1.0])
    reference -= float(reference @ axis) * axis
    if np.linalg.norm(reference) <= 1e-8:
        reference = np.array([1.0, 0.0, 0.0])
        reference -= float(reference @ axis) * axis
    reference /= np.linalg.norm(reference)
    wrapped = float(
        np.arctan2(axis @ np.cross(reference, radial), reference @ radial)
    )
    # Unwrap only against the previous accepted diagnostic value.
    delta = (wrapped - previous + np.pi) % (2.0 * np.pi) - np.pi
    return previous + float(delta), False


def execute_h1_retargeting(trial: Mapping[str, Any]) -> TrialPayload:
    """Run one paired H1 trajectory through SEW or one formal baseline."""
    models = load_models(add_simulated_tcp=True)
    baselines, sew = build_canonical_sew_target_baselines(models)
    methods = {**baselines, sew.name: sew}
    method_id = _method_id(trial)
    if method_id not in methods:
        raise ValueError(f"unknown H1 method {method_id!r}")
    method = methods[method_id]

    lower, upper = finite_joint_limits(models.master, MASTER_JOINT_NAMES)
    q_master = _master_trajectory(trial, lower=lower, upper=upper)
    sample_count = q_master.shape[0]
    q_slave_log = np.empty((sample_count, models.slave.nq))
    position_error = np.empty(sample_count)
    orientation_error = np.empty(sample_count)
    success = np.empty(sample_count, dtype=np.int8)
    smooth = np.empty(sample_count, dtype=np.int8)
    failure_code = np.empty(sample_count, dtype=np.int16)
    solver_status = np.empty(sample_count, dtype=np.int16)
    iterations = np.empty(sample_count, dtype=np.int32)
    runtime_s = np.empty(sample_count)
    solver_cost = np.empty(sample_count)
    swivel = np.empty(sample_count)
    degeneracy = np.zeros(sample_count, dtype=np.int8)
    events: list[dict[str, Any]] = []

    slave_data = models.slave.createData()
    shoulder_id = require_frame(models.slave, SLAVE_FRAMES["shoulder"])
    elbow_id = require_frame(models.slave, SLAVE_FRAMES["elbow"])
    wrist_id = require_frame(models.slave, SLAVE_FRAMES["wrist"])
    seed = None
    previous_swivel = 0.0
    for index, q_m in enumerate(q_master):
        result = method.retarget(q_m, q_slave_seed=seed)
        q_slave_log[index] = result.q_slave
        position_error[index] = result.diagnostics.position_error_m
        orientation_error[index] = result.diagnostics.orientation_error_rad
        success[index] = int(result.success)
        smooth[index] = int(result.smooth)
        failure_code[index] = _enum_code(result.failure)
        solver_status[index] = _enum_code(result.diagnostics.status)
        iterations[index] = result.diagnostics.iterations
        runtime_s[index] = result.diagnostics.runtime_s
        solver_cost[index] = result.diagnostics.cost
        previous_swivel, is_degenerate = _slave_swivel(
            models.slave,
            slave_data,
            result.q_slave,
            shoulder_id,
            elbow_id,
            wrist_id,
            previous_swivel,
        )
        swivel[index] = previous_swivel
        degeneracy[index] = int(is_degenerate)
        if result.success:
            seed = result.q_slave.copy()
        for event in result.events:
            events.append(
                {
                    "sample_index": index,
                    "event": str(event),
                    "failure_code": int(failure_code[index]),
                }
            )

    master_step = np.zeros(sample_count)
    if sample_count > 1:
        master_step[1:] = np.linalg.norm(
            (q_master[1:] - q_master[:-1] + np.pi) % (2.0 * np.pi) - np.pi,
            axis=1,
        )
    samples = {
        "sample_index": np.arange(sample_count, dtype=np.int64),
        "q_master": q_master,
        "map_q_slave": q_slave_log,
        "map_pose_success": success,
        # H1 requires a valid/smooth branch.  Differential A is evaluated in a
        # separate diagnostic study and is not silently imputed here.
        "map_differential_valid": smooth,
        "map_position_error_m": position_error,
        "map_orientation_error_rad": orientation_error,
        "map_swivel_angle_rad": swivel,
        "map_master_step_norm": master_step,
        "map_accepted": np.ones(sample_count, dtype=np.int8),
        "map_commanded_reset": np.zeros(sample_count, dtype=np.int8),
        "map_degeneracy_transition": degeneracy,
        "map_failure_code": failure_code,
        "map_solver_status": solver_status,
        "map_solver_iterations": iterations,
        "map_runtime_s": runtime_s,
        "map_solver_cost": solver_cost,
    }
    return TrialPayload(
        samples=samples,
        events=events,
        metadata={
            "evidence_scope": "pre-prototype numerical retargeting only",
            "method_id": method_id,
            "failure_enum": {
                member.value: _enum_code(member) for member in RetargetingFailure
            },
            "trajectory_family": _trajectory_spec(trial).get("family", "nominal"),
        },
    )


def _orthogonal(rng: np.random.Generator, size: int) -> np.ndarray:
    q, r = np.linalg.qr(rng.normal(size=(size, size)))
    signs = np.where(np.diag(r) >= 0.0, 1.0, -1.0)
    return q * signs


def execute_h2_synthetic(trial: Mapping[str, Any]) -> TrialPayload:
    """Run a paired, truth/estimator-separated H2 sensitivity trial."""
    method_id = _method_id(trial)
    valid_methods = {"scaled_dls", "undamped_svd", "no_bias", "no_friction"}
    if method_id not in valid_methods:
        raise ValueError(f"unknown H2 method {method_id!r}")
    trajectory = _trajectory_spec(trial)
    count = int(trajectory.get("sample_count", 256))
    if count < 8:
        raise ValueError("H2 synthetic trial needs at least eight samples")
    characteristic_length = float(
        _factor(trial, "characteristic_length_m", 0.30)
    )
    damping = float(_factor(trial, "damping", 0.02))
    min_singular = float(_factor(trial, "min_scaled_singular", 0.05))
    noise_std = float(_factor(trial, "torque_noise_std_Nm", 0.01))
    model_error_std = float(_factor(trial, "model_error_std", 0.01))
    if min_singular < 0.0 or noise_std < 0.0 or model_error_std < 0.0:
        raise ValueError("H2 perturbation factors must be non-negative")

    model_rng = generator_from_record(trial["seeds"], "model")
    sensor_rng = generator_from_record(trial["seeds"], "sensor")
    trajectory_rng = generator_from_record(trial["seeds"], "trajectory")
    u = _orthogonal(model_rng, 6)
    v = _orthogonal(model_rng, 7)
    singular = np.array([1.6, 1.25, 0.95, 0.65, 0.35, min_singular])
    base_scaled_truth = u @ np.diag(singular) @ v[:6, :]
    inverse_scaling = np.diag(
        [characteristic_length] * 3 + [1.0] * 3
    )

    phase = np.linspace(0.0, 2.0 * np.pi, count, endpoint=False)
    amplitudes = np.array([18.0, 12.0, 9.0, 1.8, 1.2, 0.8])
    offsets = trajectory_rng.uniform(-np.pi, np.pi, 6)
    wrench_reference = amplitudes[None, :] * np.sin(
        phase[:, None] * np.arange(1, 7)[None, :] + offsets[None, :]
    )
    qd = 0.6 * np.sin(
        phase[:, None] * np.arange(1, 8)[None, :]
        + trajectory_rng.uniform(-np.pi, np.pi, (1, 7))
    )
    bias = np.array([0.08, -0.05, 0.035, -0.025, 0.015, -0.01, 0.02])
    friction = JointFrictionCalibration(
        coulomb_nm=np.array([0.06, 0.05, 0.045, 0.04, 0.02, 0.02, 0.015]),
        viscous_nm_per_rad_s=np.array(
            [0.018, 0.017, 0.015, 0.014, 0.009, 0.008, 0.007]
        ),
    )
    calibration = WrenchEstimatorCalibration(
        joint_bias_nm=bias,
        friction=friction,
        characteristic_length_m=characteristic_length,
        damping=damping,
        calibration_id="g0c-synthetic-frozen-v1",
    )
    solver = (
        UndampedSVDSolver(characteristic_length)
        if method_id == "undamped_svd"
        else ScaledDLSSolver(characteristic_length, damping)
    )
    estimator = CalibratedResidualWrenchEstimator(calibration, solver)
    ablation = (
        ResidualAblation.no_bias()
        if method_id == "no_bias"
        else ResidualAblation.no_friction()
        if method_id == "no_friction"
        else ResidualAblation()
    )

    wrench_estimated = np.empty((count, 6))
    singular_log = np.empty((count, 6))
    rank = np.empty(count, dtype=np.int16)
    status = np.empty(count, dtype=np.int16)
    truth_jacobian = np.empty((count, 42))
    estimator_jacobian = np.empty((count, 42))
    residual_raw = np.empty((count, 7))
    residual_corrected = np.empty((count, 7))
    for index in range(count):
        smooth_change = 0.015 * np.sin(phase[index])
        J_truth = inverse_scaling @ (
            base_scaled_truth
            + smooth_change * model_rng.normal(size=(6, 7))
        )
        J_estimator = J_truth + inverse_scaling @ (
            model_error_std * model_rng.normal(size=(6, 7))
        )
        interaction = J_truth.T @ wrench_reference[index]
        measured = (
            interaction
            + bias
            + friction.torque(qd[index])
            + sensor_rng.normal(0.0, noise_std, 7)
        )
        estimate = estimator.estimate(
            J_estimator,
            measured_torque_nm=measured,
            rigid_body_torque_nm=np.zeros(7),
            joint_velocity_rad_s=qd[index],
            ablation=ablation,
        )
        wrench_estimated[index] = estimate.solve.wrench
        singular_log[index] = estimate.solve.singular_values
        rank[index] = estimate.solve.rank
        status[index] = _enum_code(estimate.solve.status)
        truth_jacobian[index] = J_truth.reshape(-1)
        estimator_jacobian[index] = J_estimator.reshape(-1)
        residual_raw[index] = estimate.residual.raw_residual_nm
        residual_corrected[index] = estimate.residual.residual_nm

    return TrialPayload(
        samples={
            "sample_index": np.arange(count, dtype=np.int64),
            "wrench_reference": wrench_reference,
            "wrench_estimated": wrench_estimated,
            "wrench_sample_mask": np.ones(count, dtype=np.int8),
            "qd_slave": qd,
            "jacobian_truth": truth_jacobian,
            "jacobian_estimator": estimator_jacobian,
            "scaled_singular_values": singular_log,
            "solver_rank": rank,
            "solver_status": status,
            "tau_residual_raw": residual_raw,
            "tau_residual_corrected": residual_corrected,
        },
        metadata={
            "evidence_scope": (
                "synthetic sensitivity only; not independent physical F/T evidence"
            ),
            "method_id": method_id,
            "truth_estimator_models_separated": True,
            "calibration_id": calibration.calibration_id,
        },
    )


def execute_bilateral_simulation(trial: Mapping[str, Any]) -> TrialPayload:
    """Run one paired H3/H4 rigid-body trial with a frozen network trace."""
    method_id = _method_id(trial)
    scenarios = {scenario.key: scenario for scenario in SCENARIOS}
    if method_id not in scenarios:
        raise ValueError(f"unknown bilateral method {method_id!r}")
    scenario = scenarios[method_id]
    map_policy = str(_factor(trial, "map_policy", scenario.map_policy))
    scenario = replace(scenario, map_policy=map_policy)

    trajectory = _trajectory_spec(trial)
    trajectory_rng = generator_from_record(trial["seeds"], "trajectory")
    seed = int(trajectory_rng.integers(0, np.iinfo(np.int32).max))
    config = replace(
        SimulationConfig(),
        seed=seed,
        duration=float(trajectory.get("duration_s", 1.2)),
        contact_probe_fraction=float(
            trajectory.get("contact_probe_fraction", 0.0)
        ),
        feedback_delay_s=float(_factor(trial, "return_delay_s", 0.04)),
        forward_delay_s=float(_factor(trial, "forward_delay_s", 0.0)),
        return_jitter_s=float(_factor(trial, "return_jitter_s", 0.0)),
        forward_jitter_s=float(_factor(trial, "forward_jitter_s", 0.0)),
        return_packet_loss=float(_factor(trial, "return_packet_loss", 0.0)),
        forward_packet_loss=float(_factor(trial, "forward_packet_loss", 0.0)),
        wall_stiffness=float(_factor(trial, "wall_stiffness", 800.0)),
        wall_damping=float(_factor(trial, "wall_damping", 45.0)),
    )
    models = load_models(add_simulated_tcp=True)
    mapper = build_mapper(models)
    wall, wall_metadata, q_slave_start = make_wall(config, models, mapper)
    if trajectory.get("family") == "free_space":
        travel = float(wall_metadata["free_space_travel_m"])
        wall = replace(
            wall,
            point=wall.point + 2.0 * travel * wall.normal,
        )
        wall_metadata = {
            **wall_metadata,
            "condition": "free_space",
            "point_world_m": wall.point.tolist(),
        }
    result = simulate_scenario(
        scenario, config, models, wall, q_slave_start
    )
    samples = {name: value.copy() for name, value in result.logs.items()}
    samples["tau_master_raw"] = samples["tau_master_mapped"].copy()
    samples["tau_slave_source"] = (
        samples["tau_slave_residual_source"].copy()
        if scenario.mapping == "differential_residual"
        else samples["tau_slave_matched_wrench"].copy()
    )
    samples["return_valid"] = samples["return_packet_active"].astype(np.int8)
    samples["energy_before_J"] = samples["energy_before"].copy()
    samples["energy_after_J"] = samples["tank_energy"].copy()
    samples["energy_preclip_J"] = samples["energy_preclip"].copy()
    return TrialPayload(
        samples=samples,
        events=(),
        metadata={
            "evidence_scope": (
                "pre-prototype rigid-body simulation; no physical or human claim"
            ),
            "scenario": {
                "key": scenario.key,
                "mapping": scenario.mapping,
                "supervisor": scenario.supervisor,
                "map_policy": scenario.map_policy,
            },
            "wall": wall_metadata,
            "online_metrics_are_diagnostic_only": result.metrics,
        },
    )