exoskeleton/code/test/ablation_haptic.py

# -*- coding: utf-8 -*-
"""
ablation_haptic.py

Quantitative ablation for the haptic rendering module.

Modes:
- baseline   : full method (energy tank + F/V filtering + alpha smoothing + torque LPF)
- no_tank    : energy tank disabled (no passivity enforcement)
- no_filter  : tank enabled but F/V filtering + alpha smoothing + torque LPF disabled

Run:
    python test/ablation_haptic.py
from the project root, with config/config.yaml following the same
structure as demo.py and test_sew.py.
"""

import os
import sys
import time
from typing import Dict

import numpy as np
import pinocchio as pin

import matplotlib.pyplot as plt
from omegaconf import OmegaConf

# Make sure project root is on sys.path
THIS_DIR = os.path.dirname(__file__)
PROJECT_ROOT = os.path.abspath(os.path.join(THIS_DIR, ".."))
if PROJECT_ROOT not in sys.path:
    sys.path.append(PROJECT_ROOT)

from core.sew_mapper import SEWMapper
from core.interaction_estimater import InteractionEstimator
from core.haptic_render import HapticRenderer


# -------------------- Utilities --------------------

def dummy_master_measure(step: int, model: pin.Model) -> np.ndarray:
    """
    Synthetic master joint measurement used for offline ablation.
    Start from neutral and add small sinusoidal motions on the first few joints.
    """
    q = pin.neutral(model)
    n_use = min(model.nq, 6)
    for i in range(n_use):
        q[i] = 0.5 * np.sin(0.01 * step + 0.7 * i)
    return q


def dummy_slave_dynamics(q_cmd: np.ndarray, model: pin.Model) -> np.ndarray:
    """
    Simple slave "execution" model: follow the commanded joint positions
    with a small Gaussian disturbance, to emulate interaction / noise.
    """
    q_cmd = np.asarray(q_cmd, dtype=float).reshape(-1,)
    noise = 0.001 * np.random.randn(model.nq)
    return q_cmd + noise


def load_config():
    """
    Load config/config.yaml in the project root (same style as demo.py / test_sew.py).
    """
    cand = os.path.join(PROJECT_ROOT, "config", "config.yaml")
    if not os.path.exists(cand):
        raise FileNotFoundError(f"config.yaml not found at: {cand}")
    conf = OmegaConf.load(cand)
    print(f"[INFO] Loaded config from: {cand}")
    return conf


def build_modules(conf):
    # Build models from URDF
    master_model = pin.buildModelFromUrdf(str(conf.master_urdf))
    slave_model = pin.buildModelFromUrdf(str(conf.slave_urdf))

    # SEW mapper
    sew = SEWMapper(
        master_model=master_model,
        slave_model=slave_model,
        m_shoulder_frame=conf.m_shoulder_frame,
        m_elbow_frame=conf.m_elbow_frame,
        m_wrist_frame=conf.m_wrist_frame,
        m_ee_frame=conf.m_ee_frame,
        s_shoulder_frame=conf.s_shoulder_frame,
        s_elbow_frame=conf.s_elbow_frame,
        s_wrist_frame=conf.s_wrist_frame,
        s_ee_frame=conf.s_ee_frame,
        slave_joint_names=conf.sew_mapper.slave_joint_names,
        up_dir=np.array(conf.sew_mapper.up_dir, dtype=float),
        eps_clip=float(conf.sew_mapper.eps_clip),
    )

    # Interaction estimator on slave side
    estimator = InteractionEstimator(
        model=slave_model,
        chest_frame_name=conf.s_base_frame,
        ee_frame_name=conf.s_ee_frame,
        lambda_damp=float(conf.interaction_est.lambda_damp),
    )

    return master_model, slave_model, sew, estimator



# -------------------- Haptic renderer variants --------------------

class HapticRendererNoTank(HapticRenderer):
    """
    Variant that completely disables the energy tank:
    - directly maps interaction wrench to master joint torques through CJ_m^T
    - no passivity enforcement
    """

    def _compute_feedback_tau(self,
                              q_m: np.ndarray,
                              qd_m: np.ndarray,
                              CF_int_slave_C: np.ndarray,
                              V_slave_C: np.ndarray | None,
                              dt: float):
        CF = np.asarray(CF_int_slave_C, dtype=float).reshape(6,)
        # Chest Jacobian on master side
        CJ_m = self.CJ_master.chest_jacobian(q_m, qd_m)
        tau_fb_m_raw = CJ_m.T @ CF
        alpha = 1.0
        return tau_fb_m_raw, alpha


def make_renderer(conf, master_model: pin.Model, mode: str) -> HapticRenderer:
    """
    Construct different haptic renderers for ablation.
    """
    common_kwargs = dict(
        master_model=master_model,
        chest_frame_name=conf.m_base_frame,
        ee_frame_name=conf.m_ee_frame,
        feedback_strength=float(conf.haptic_render.feedback_strength),
        E_init=float(conf.haptic_render.E_init),
        E_max=float(conf.haptic_render.E_max),
        alpha_floor=float(conf.haptic_render.alpha_floor),
        alpha_ceil=float(conf.haptic_render.alpha_ceil),
        E0=float(conf.haptic_render.E0),
    )

    if mode == "baseline":
        renderer = HapticRenderer(**common_kwargs)

        # baseline：能量罐 + 强滤波 + 平滑 + 扭矩低通（“管得很严”）
        renderer.tank.tp.force_alpha = 0.05
        renderer.tank.tp.vel_alpha = 0.05
        renderer.tank.tp.alpha_smooth = 0.02
        renderer.tank.tp.power_deadzone = 0.0
        renderer.tau_alpha = 0.05

    elif mode == "no_tank":
        renderer = HapticRendererNoTank(**common_kwargs)

        # no_tank：完全不做能量控制，也不做扭矩低通，暴露最差情况
        renderer.tau_alpha = 1.0

    elif mode == "no_filter":
        renderer = HapticRenderer(**common_kwargs)

        # no_filter：保留能量罐，但不给它任何滤波/平滑能力
        renderer.tank.tp.force_alpha = 1.0
        renderer.tank.tp.vel_alpha = 1.0
        renderer.tank.tp.alpha_smooth = 1.0
        renderer.tank.tp.power_deadzone = 0.0
        renderer.tau_alpha = 1.0

    else:
        raise ValueError(f"Unknown ablation mode: {mode}")

    return renderer


# -------------------- Teleoperation simulation (offline) --------------------

def simulate_teleop(conf,
                    master_model: pin.Model,
                    slave_model: pin.Model,
                    sew: SEWMapper,
                    estimator: InteractionEstimator,
                    renderer: HapticRenderer,
                    Ts: float = 0.002,
                    steps: int = 800) -> Dict[str, np.ndarray]:
    """
    Offline simulation similar to demo.py, but without real-time delays.
    Returns logs for computing quantitative metrics.
    """

    # States
    q_slave = pin.neutral(slave_model)
    dq_slave = np.zeros(slave_model.nv)
    q_slave_prev = q_slave.copy()

    q_m_init = dummy_master_measure(0, master_model)
    pre_q_m = q_m_init.copy()
    pre_qd_m = np.zeros_like(q_m_init)

    # Logs
    log_tau_fb = []
    log_alpha = []
    log_E = []
    log_tau_int_norm = []
    log_qd_m = []
    ee_traj_world = []

    # Harder virtual wall to highlight differences
    d_wall = 0.35
    k_wall = 20000.0

    for k in range(steps):
        # --- master state ---
        q_m = dummy_master_measure(k, master_model)
        if k == 0:
            qd_m = np.zeros_like(q_m)
            qdd_m = np.zeros_like(q_m)
        else:
            qd_m = (q_m - pre_q_m) / Ts
            qdd_m = (qd_m - pre_qd_m) / Ts

        pre_q_m = q_m
        pre_qd_m = qd_m
        log_qd_m.append(qd_m.copy())

        # --- SEW retargetting: master -> desired slave ---
        q_des_slave, _ = sew.retargetting(q_m, q_slave)

        # --- slave motion (with small noise) ---
        q_slave = dummy_slave_dynamics(q_des_slave, slave_model)
        dq_slave = (q_slave - q_slave_prev) / Ts
        q_slave_prev = q_slave.copy()
        qdd_slave = np.zeros(slave_model.nv)

        # --- virtual wall on slave side ---
        CJ_env = estimator._chest_jacobian(q_slave, dq_slave)

        oTC = estimator.data.oMf[estimator.fid_C]   # ^wT_C
        oTEE = estimator.data.oMf[estimator.fid_EE] # ^wT_EE

        R_wc = oTC.rotation
        p_wc = oTC.translation
        R_cw = R_wc.T
        p_cw = -R_cw @ p_wc

        p_we = oTEE.translation
        ee_traj_world.append(p_we.copy())
        p_ce = R_cw @ p_we + p_cw
        x_ce = float(p_ce[0])

        if x_ce > d_wall:
            delta = x_ce - d_wall
            Fx = -k_wall * delta
            CF_env = np.array([Fx, 0, 0, 0, 0, 0], dtype=float)
        else:
            CF_env = np.zeros(6, dtype=float)

        tau_env = CJ_env.T @ CF_env

        # --- interaction estimation on slave ---
        tau_model = estimator._tau_model(q_slave, dq_slave, qdd_slave)
        tau_meas = tau_model + tau_env
        tau_int, CF_int, CJ = estimator.estimate(q_slave, dq_slave, qdd_slave, tau_meas)

        log_tau_int_norm.append(np.linalg.norm(tau_int))

        V_slave_C = CJ @ dq_slave  # 6 x 1

        # --- haptic rendering on master side ---
        tau_cmd_m, tau_fb_m, alpha = renderer.render_tau(
            q_m=q_m,
            qd_m=qd_m,
            qdd_m=qdd_m,
            CF_int_slave_C=CF_int,
            V_slave_C=V_slave_C,
            tau_ff_fric=None,
            dt=Ts,
        )

        log_tau_fb.append(tau_fb_m.copy())
        log_alpha.append(alpha)
        # not all variants really use the tank, but querying E is safe
        log_E.append(getattr(renderer.tank, "E", 0.0))

    # stack
    log_tau_fb = np.vstack(log_tau_fb)
    log_alpha = np.asarray(log_alpha)
    log_E = np.asarray(log_E)
    log_tau_int_norm = np.asarray(log_tau_int_norm)
    log_qd_m = np.vstack(log_qd_m)
    ee_traj_world = np.vstack(ee_traj_world)

    return dict(
        tau_fb=log_tau_fb,
        alpha=log_alpha,
        E=log_E,
        tau_int_norm=log_tau_int_norm,
        qd_m=log_qd_m,
        ee_traj_world=ee_traj_world,
    )


# -------------------- Metric computation --------------------

def compute_metrics(tau_fb: np.ndarray,
                    qd_m: np.ndarray,
                    E: np.ndarray,
                    Ts: float) -> Dict[str, float]:
    """
    Compute quantitative metrics for ablation:
    - rms / peak torque
    - torque spike count / rate
    - positive power injection ratio
    - (optional) energy range if tank is used
    """
    # torque norms
    tau_norm = np.linalg.norm(tau_fb, axis=1)
    rms_tau = float(np.sqrt(np.mean(tau_norm ** 2)))
    peak_tau = float(np.max(tau_norm))

    # torque spikes: ||Δτ||_∞ > threshold
    dtau = np.diff(tau_fb, axis=0)
    dtau_inf = np.max(np.abs(dtau), axis=1)
    spike_th = 5.0  # Nm, can be adjusted to your system scale
    n_spikes = int(np.sum(dtau_inf > spike_th))
    spike_rate = n_spikes / (len(dtau_inf) * Ts)

    # positive power injection ratio: tau_fb^T * qd_m > 0
    P = np.sum(tau_fb * qd_m, axis=1)
    P_inject_ratio = float(np.mean(P > 0.0))

    metrics = dict(
        rms_tau=rms_tau,
        peak_tau=peak_tau,
        n_spikes=n_spikes,
        spike_rate=spike_rate,
        P_inject_ratio=P_inject_ratio,
    )

    if np.any(E != 0.0):
        metrics["E_min"] = float(np.min(E))
        metrics["E_max"] = float(np.max(E))

    return metrics


def plot_ablation_results(all_logs, all_metrics, Ts):

    modes = ["baseline", "no_tank", "no_filter"]
    colors = {
        "baseline": "tab:blue",
        "no_tank": "tab:red",
        "no_filter": "tab:green",
    }

    # -------------------------------------------
    # 1) Torque feedback trajectories (norm vs time)
    # -------------------------------------------
    plt.figure(figsize=(8,4))
    for mode in modes:
        tau_norm = np.linalg.norm(all_logs[mode]["tau_fb"], axis=1)
        t = np.arange(len(tau_norm)) * Ts
        plt.plot(t, tau_norm, label=mode, color=colors[mode])
    plt.xlabel("Time [s]")
    plt.ylabel(r"$\|\tau_{\mathrm{fb}}(t)\|\ \mathrm{[N\cdot m]}$")
    plt.title("Feedback Torque Norm Over Time")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("logs_ablation/fig_tau_fb_norm.png", dpi=300)

    # -------------------------------------------
    # 2) Alpha(t) (energy tank scaling)
    # -------------------------------------------
    plt.figure(figsize=(8,4))
    for mode in modes:
        alpha = all_logs[mode]["alpha"]
        t = np.arange(len(alpha)) * Ts
        plt.plot(t, alpha, label=mode, color=colors[mode])
    plt.xlabel("Time [s]")
    plt.ylabel(r"$\alpha(t)$")
    plt.title("Energy Tank Scaling Factor")
    plt.ylim([-0.1, 1.1])
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("logs_ablation/fig_alpha.png", dpi=300)

    # -------------------------------------------
    # 3) Tank Energy E(t)
    # -------------------------------------------
    plt.figure(figsize=(8,4))
    for mode in modes:
        E = all_logs[mode]["E"]
        t = np.arange(len(E)) * Ts
        plt.plot(t, E, label=mode, color=colors[mode])
    plt.xlabel("Time [s]")
    plt.ylabel(r"$E(t)$")
    plt.title("Energy Tank Level")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("logs_ablation/fig_energy.png", dpi=300)

    # -------------------------------------------
    # 4) Bar charts: RMS / Peak / Spikes / P>0
    # -------------------------------------------
    metrics_list = ["rms_tau", "peak_tau", "n_spikes", "P_inject_ratio"]
    metric_names = [
        "RMS Torque",
        "Peak Torque",
        "# Torque Spikes",
        "Positive Power Ratio",
    ]

    plt.figure(figsize=(9,4))
    for i, key in enumerate(metrics_list):
        plt.subplot(1,4,i+1)
        vals = [all_metrics[m][key] for m in modes]
        plt.bar(modes, vals, color=[colors[m] for m in modes])
        plt.title(metric_names[i])
        plt.xticks(rotation=45)
        plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("logs_ablation/fig_metrics_barchart.png", dpi=300)

    print("[INFO] All plots saved to logs_ablation/")



# -------------------- Main --------------------

def main():
    conf = load_config()
    master_model, slave_model, sew, estimator = build_modules(conf)

    Ts = 0.002
    steps = 800

    modes = ["baseline", "no_tank", "no_filter"]
    all_logs = {}
    all_metrics = {}

    print(f"dt = {Ts:.4f} s, steps = {steps}")

    for mode in modes:
        print(f"\n========== Running haptic ablation: {mode} ==========")
        renderer = make_renderer(conf, master_model, mode)
        renderer.reset_tank(float(conf.haptic_render.E0))

        # fix random seed so that each mode sees the same slave noise
        np.random.seed(0)

        t0 = time.time()
        logs = simulate_teleop(
            conf,
            master_model,
            slave_model,
            sew,
            estimator,
            renderer,
            Ts=Ts,
            steps=steps,
        )
        elapsed = time.time() - t0

        all_logs[mode] = logs

        metrics = compute_metrics(
            tau_fb=logs["tau_fb"],
            qd_m=logs["qd_m"],
            E=logs["E"],
            Ts=Ts,
        )
        all_metrics[mode] = metrics

        print(
            f"[{mode}] done in {elapsed:.3f} s | "
            f"rms_tau={metrics['rms_tau']:.3f}, "
            f"peak_tau={metrics['peak_tau']:.3f}, "
            f"spikes={metrics['n_spikes']}, "
            f"spike_rate={metrics['spike_rate']:.3f} 1/s, "
            f"P>0={metrics['P_inject_ratio']:.3f}, "
            f"E_range="
            f"{'[{:.2f},{:.2f}]'.format(metrics['E_min'], metrics['E_max']) if 'E_min' in metrics else 'N/A'}"
        )

    # Save logs for further plotting if needed
    save_dir = os.path.join(PROJECT_ROOT, "logs_ablation")
    os.makedirs(save_dir, exist_ok=True)
    save_path = os.path.join(save_dir, "haptics_ablation_results.npz")

    # pack as a single dict and rely on pickle when loading
    np.savez(
        save_path,
        all_logs=all_logs,
        all_metrics=all_metrics,
        Ts=Ts,
    )

    print(f"\n[INFO] Ablation logs and metrics saved to:\n  {save_path}")
    plot_ablation_results(all_logs, all_metrics, Ts)


if __name__ == "__main__":
    main()