exoskeleton/code/core/haptic_render.py

# haptic_render.py
# -*- coding: utf-8 -*-
from __future__ import annotations

from dataclasses import dataclass

import numpy as np

try:
    import pinocchio as pin
except ModuleNotFoundError:  # Pure supervisor tests do not require Pinocchio.
    pin = None


# ===== Final applied-port energy supervision =====

@dataclass
class TankParams:
    E_min: float = 0.5
    E_max: float = 20.0
    alpha_floor: float = 0.0
    alpha_ceil: float = 1.0
    power_epsilon: float = 1e-12

    # Legacy fields are kept so existing configuration/ablation code can load.
    # They are deliberately not used for energy accounting: filtering, a power
    # dead-zone, or alpha smoothing at this stage would make the tank account a
    # different signal from the torque actually applied at the master port.
    force_alpha: float = 0.2
    vel_alpha: float = 0.2
    alpha_smooth: float = 0.1
    power_deadzone: float = 0.0


@dataclass(frozen=True)
class AppliedPortDiagnostics:
    """Snapshot of one final-port projection/accounting step."""

    tau_candidate: np.ndarray
    tau_applied: np.ndarray
    rho: float
    E_before: float
    E_preclip: float
    E_after: float
    candidate_power: float
    power: float
    fail_safe_active: bool


class AppliedPortEnergySupervisor:
    """Passivity supervisor at the final master haptic-torque port.

    Positive ``tau_app @ qd_m`` is power delivered by the device to the
    operator and therefore discharges the tank.  Negative power charges the
    tank.  The candidate torque passed to :meth:`apply` must
    already include all nominal scaling, filtering, rate limiting, and actuator
    saturation.  The returned torque is final: downstream code must not filter,
    rescale, rate-limit, or saturate it.

    The projection is radial, ``tau_app = alpha * tau_candidate``.  This keeps a
    previously saturated candidate inside its actuator limits while imposing
    the exact one-step energy budget.  ``alpha_floor`` is never allowed to
    override the energy-safe upper bound.
    """

    def __init__(self, tank: TankParams | None = None, E0: float = 2.0):
        self.tp = TankParams() if tank is None else tank
        if not np.isfinite(self.tp.E_min) or not np.isfinite(self.tp.E_max):
            raise ValueError("E_min and E_max must be finite")
        if self.tp.E_max < self.tp.E_min:
            raise ValueError("E_max must be greater than or equal to E_min")
        if not 0.0 <= self.tp.alpha_ceil <= 1.0:
            raise ValueError("alpha_ceil must lie in [0, 1]")
        if not 0.0 <= self.tp.alpha_floor <= self.tp.alpha_ceil:
            raise ValueError("alpha_floor must lie in [0, alpha_ceil]")
        if self.tp.power_epsilon < 0.0:
            raise ValueError("power_epsilon must be non-negative")

        self.E = float(np.clip(E0, self.tp.E_min, self.tp.E_max))
        self.alpha = 1.0
        self.last_alpha = 1.0
        self.last_power = 0.0
        self.last_energy_before = self.E
        self.last_energy_after = self.E
        self.last_tau_candidate = np.empty(0, dtype=float)
        self.last_tau_applied = np.empty(0, dtype=float)
        self.fail_safe_active = False
        self.last_diagnostics: AppliedPortDiagnostics | None = None

    def reset(self, E0: float | None = None):
        if E0 is None:
            E0 = self.tp.E_min
        self.E = float(np.clip(E0, self.tp.E_min, self.tp.E_max))
        self.alpha = 1.0
        self.last_alpha = 1.0
        self.last_power = 0.0
        self.last_energy_before = self.E
        self.last_energy_after = self.E
        self.last_tau_candidate = np.empty(0, dtype=float)
        self.last_tau_applied = np.empty(0, dtype=float)
        self.fail_safe_active = False
        self.last_diagnostics = None

    def _record(self,
                tau_candidate: np.ndarray,
                tau_applied: np.ndarray,
                rho: float,
                E_before: float,
                E_preclip: float,
                E_after: float,
                candidate_power: float,
                power: float,
                fail_safe_active: bool) -> AppliedPortDiagnostics:
        """Record one result while retaining legacy scalar/state attributes."""
        diagnostics = AppliedPortDiagnostics(
            tau_candidate=tau_candidate.copy(),
            tau_applied=tau_applied.copy(),
            rho=float(rho),
            E_before=float(E_before),
            E_preclip=float(E_preclip),
            E_after=float(E_after),
            candidate_power=float(candidate_power),
            power=float(power),
            fail_safe_active=bool(fail_safe_active),
        )
        self.alpha = diagnostics.rho
        self.last_alpha = diagnostics.rho
        self.last_power = diagnostics.power
        self.last_energy_before = diagnostics.E_before
        self.last_energy_after = diagnostics.E_after
        self.last_tau_candidate = diagnostics.tau_candidate.copy()
        self.last_tau_applied = diagnostics.tau_applied.copy()
        self.fail_safe_active = diagnostics.fail_safe_active
        self.last_diagnostics = diagnostics
        return diagnostics

    def _zero_fail_safe(
        self,
        tau_candidate: np.ndarray,
    ) -> tuple[np.ndarray, AppliedPortDiagnostics]:
        tau_zero = np.zeros_like(tau_candidate, dtype=float)
        diagnostics = self._record(
            tau_candidate=tau_candidate,
            tau_applied=tau_zero,
            rho=0.0,
            E_before=self.E,
            E_preclip=self.E,
            E_after=self.E,
            candidate_power=0.0,
            power=0.0,
            fail_safe_active=True,
        )
        return tau_zero, diagnostics

    def apply(self,
              tau_raw: np.ndarray,
              qd: np.ndarray,
              dt: float) -> tuple[np.ndarray, AppliedPortDiagnostics]:
        """Return the final safe torque and diagnostics, accounting exactly once.

        ``tau_raw`` is the already shaped and actuator-limited candidate at the
        master joint port.  It may come directly from a mapping such as
        ``A.T @ tau_s`` or ``J_m.T @ F``.  ``qd`` is the velocity at that same
        port.  The returned ``tau_app`` must be sent unchanged to the haptic
        output (apart from an emergency zero-output fail-safe).
        """
        tau_candidate = np.asarray(tau_raw, dtype=float).reshape(-1,)
        qd_m = np.asarray(qd, dtype=float).reshape(-1,)
        if tau_candidate.shape != qd_m.shape:
            raise ValueError(
                "tau_raw and qd must have identical one-dimensional shapes"
            )
        if (
            tau_candidate.size == 0
            or not np.isfinite(dt)
            or dt <= 0.0
            or not np.all(np.isfinite(tau_candidate))
            or not np.all(np.isfinite(qd_m))
        ):
            return self._zero_fail_safe(tau_candidate)

        energy_before = self.E
        nominal_power = float(np.dot(tau_candidate, qd_m))
        if not np.isfinite(nominal_power):
            return self._zero_fail_safe(tau_candidate)
        rho = self.tp.alpha_ceil

        # Only device-to-human power consumes stored energy.  alpha_floor is a
        # preference, not a safety constraint, and cannot force an unsafe gain.
        if nominal_power > 0.0:
            available = max(0.0, energy_before - self.tp.E_min)
            rho_energy = available / (nominal_power * dt)
            rho = min(rho, max(0.0, rho_energy))

        tau_app = rho * tau_candidate
        applied_power = float(np.dot(tau_app, qd_m))

        # This is the sole tank update.  It uses the exact torque returned to
        # the caller and the exact master velocity supplied for this sample.
        energy_preclip = energy_before - applied_power * dt
        self.E = float(np.clip(energy_preclip, self.tp.E_min, self.tp.E_max))
        diagnostics = self._record(
            tau_candidate=tau_candidate,
            tau_applied=tau_app,
            rho=rho,
            E_before=energy_before,
            E_preclip=energy_preclip,
            E_after=self.E,
            candidate_power=nominal_power,
            power=applied_power,
            fail_safe_active=False,
        )
        return tau_app, diagnostics

    def project_and_account(self,
                            tau_candidate: np.ndarray,
                            qd_m: np.ndarray,
                            dt: float) -> tuple[float, np.ndarray]:
        """Compatibility wrapper around :meth:`apply`.

        Returns ``(alpha, tau_app)``.  The wrapper invokes ``apply`` exactly
        once, so it cannot introduce a second tank update.
        """
        tau_app, diagnostics = self.apply(tau_candidate, qd_m, dt)
        return diagnostics.rho, tau_app

    def observe_master(self, *_args, **_kwargs):
        """Reject the former second accounting path."""
        raise RuntimeError(
            "observe_master() is disabled: use apply() exactly once "
            "at the final applied master-torque port"
        )

    def enforce(self, *_args, **_kwargs):
        """Reject wrench/twist-port enforcement from the former implementation."""
        raise RuntimeError(
            "enforce(CF, VC, dt) is disabled: passivity is enforced on final "
            "master tau_app and qd_m via apply()"
        )


# Backward-compatible class name for existing imports.
EnergyTankPOPC = AppliedPortEnergySupervisor


# ===== 主端胸腔雅可比 =====

class ChestJacobianMaster:
    def __init__(self, model: pin.Model,
                 chest_frame_name: str,
                 ee_frame_name: str):
        if pin is None:
            raise ModuleNotFoundError(
                "Pinocchio is required to construct ChestJacobianMaster"
            )
        self.model = model
        self.data  = model.createData()
        self.fid_C  = model.getFrameId(chest_frame_name)
        self.fid_EE = model.getFrameId(ee_frame_name)

    def chest_jacobian(self, q_m, qd_m):
        pin.forwardKinematics(self.model, self.data, q_m, qd_m)
        pin.updateFramePlacements(self.model, self.data)

        # LOCAL_WORLD_ALIGNED expresses the twist at the EE point in axes
        # parallel to the world frame.  The spatial-vector convention used by
        # this project is [linear; angular].
        Jw = pin.computeFrameJacobian(self.model, self.data,
                                      q_m, self.fid_EE,
                                      pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)

        # Rotate both vector blocks into chest axes, but do not translate the
        # reference point: the wrench paired with this Jacobian still acts at
        # the EE point.  A full spatial adjoint here would introduce a spurious
        # moment-arm term between the chest origin and EE point.
        oTC = self.data.oMf[self.fid_C]  # ^wT_C
        R_cw = oTC.rotation.T
        rotate_axes = np.zeros((6, 6))
        rotate_axes[:3, :3] = R_cw
        rotate_axes[3:, 3:] = R_cw
        CJ_m = rotate_axes @ Jw
        return CJ_m


# ===== 主端力反馈渲染 =====

class HapticRenderer:
    """Render interaction wrench as final, passivity-supervised master torque.

    The haptic path has one strict ordering:

    ``J_m.T @ F`` (direct baseline)
      -> optional nominal low-pass
      -> nominal sign/strength scaling
      -> torque-rate limiting
      -> actuator torque saturation
      -> final applied-port energy projection and single accounting update.

    No filtering or limiting is performed after the energy projection.  The
    second returned torque is therefore the exact ``tau_app`` used in
    ``tau_app @ qd_m`` for the tank update.
    """

    def __init__(self,
                 master_model: pin.Model,
                 chest_frame_name: str,
                 ee_frame_name: str,
                 feedback_strength: float,
                 E_init: float,
                 E_max: float,
                 alpha_floor: float,
                 alpha_ceil: float,
                 E0: float = 2.0,
                 torque_limit: float | np.ndarray | None = None,
                 torque_rate_limit: float | np.ndarray | None = None,
                 tau_filter_alpha: float = 0.2):
        # 约定 feedback_strength > 0，反作用时乘上一个负号
        self.feedback_strength = float(feedback_strength)
        if not np.isfinite(self.feedback_strength):
            raise ValueError("feedback_strength must be finite")

        self.CJ_master = ChestJacobianMaster(master_model,
                                             chest_frame_name,
                                             ee_frame_name)

        tank_params = TankParams(
            E_min=E_init,
            E_max=E_max,
            alpha_floor=alpha_floor,
            alpha_ceil=alpha_ceil,
        )
        self.tank = AppliedPortEnergySupervisor(tank_params, E0)

        # Nominal shaping state.  All of this is upstream of the energy
        # projection.  Public tau_alpha is retained for existing scripts.
        self.tau_alpha = float(tau_filter_alpha)
        if not 0.0 <= self.tau_alpha <= 1.0:
            raise ValueError("tau_filter_alpha must lie in [0, 1]")
        self.torque_limit = torque_limit
        self.torque_rate_limit = torque_rate_limit
        self._tau_fb_state = np.zeros(master_model.nv)
        self._tau_applied_prev = np.zeros(master_model.nv)
        self.last_rate_limit_active = False
        self.last_torque_saturation_active = False

    def reset_tank(self, E0: float = None):
        self.tank.reset(E0)
        self._tau_fb_state[:] = 0.0
        self._tau_applied_prev[:] = 0.0
        self.last_rate_limit_active = False
        self.last_torque_saturation_active = False

    # ---------- 主端逆动力学：计算 tau_master_current ----------

    def _inverse_dynamics(self,
                          q_m: np.ndarray,
                          qd_m: np.ndarray,
                          qdd_m: np.ndarray,
                          tau_ff_fric: np.ndarray | None = None) -> np.ndarray:
        """
        主端动力学模型：
            τ = M(q) qdd + C(q,qd) qd + g(q) + τ_ff_fric
        """
        q_m  = np.asarray(q_m,  dtype=float).reshape(-1,)
        qd_m = np.asarray(qd_m, dtype=float).reshape(-1,)
        qdd_m = np.asarray(qdd_m, dtype=float).reshape(-1,)
        if tau_ff_fric is not None:
            tau_ff_fric = np.asarray(tau_ff_fric, dtype=float).reshape(-1,)

        model = self.CJ_master.model
        data  = self.CJ_master.data

        # 惯量矩阵 M(q)
        M = pin.crba(model, data, q_m)
        M = (M + M.T) - np.diag(M.diagonal())   # 数值对称化

        # 非线性项（科氏/离心 + 重力）
        nle = pin.nonLinearEffects(model, data, q_m, qd_m)  # = C(q,qd)qd + g(q)
        g   = pin.computeGeneralizedGravity(model, data, q_m)
        Cqd = nle - g

        tau = M @ qdd_m + Cqd + g
        if tau_ff_fric is not None:
            tau = tau + tau_ff_fric
        return tau

    def compute_tau_master_current(self,
                                   q_m: np.ndarray,
                                   qd_m: np.ndarray,
                                   qdd_m: np.ndarray,
                                   tau_ff_fric: np.ndarray | None = None) -> np.ndarray:
        """
        对外接口：给定主端 (q, qd, qdd)，计算
            τ_master_current = M qdd + C qd + g + τ_ff_fric
        """
        return self._inverse_dynamics(q_m, qd_m, qdd_m, tau_ff_fric)

    # ---------- 从端交互力 → 主端反馈扭矩 ----------

    @staticmethod
    def _symmetric_limit_vector(value: float | np.ndarray | None,
                                size: int,
                                name: str) -> np.ndarray:
        if value is None:
            return np.full(size, np.inf, dtype=float)
        limit = np.asarray(value, dtype=float)
        if limit.ndim == 0:
            limit = np.full(size, float(limit), dtype=float)
        else:
            limit = limit.reshape(-1,)
            if limit.size != size:
                raise ValueError(f"{name} must be scalar or have {size} entries")
        if np.any(np.isnan(limit)) or np.any(limit < 0.0):
            raise ValueError(f"{name} must contain non-negative values")
        return limit

    def direct_from_CF(self,
                       q_m: np.ndarray,
                       qd_m: np.ndarray,
                       CF_int_slave_C: np.ndarray) -> np.ndarray:
        """Return the unmodified baseline mapping ``J_m.T @ F``.

        This interface intentionally applies no sign convention, nominal gain,
        filter, rate limit, saturation, or energy supervision.
        """
        CF = np.asarray(CF_int_slave_C, dtype=float).reshape(6,)
        CJ_m = self.CJ_master.chest_jacobian(q_m, qd_m)
        return np.asarray(CJ_m.T @ CF, dtype=float).reshape(-1,)

    # Explicitly named alias for experiment/baseline code.
    render_direct_from_CF = direct_from_CF

    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):
        """Legacy extension hook.

        The base implementation returns the direct mapping and ``None`` to
        request final applied-port supervision.  Existing no-tank subclasses
        that override this method and return a numeric alpha continue to bypass
        only the energy projection; all upstream nominal actuator shaping still
        applies.  ``V_slave_C`` is retained solely for call compatibility and is
        never used for tank accounting.
        """
        del V_slave_C, dt
        return self.direct_from_CF(q_m, qd_m, CF_int_slave_C), None

    def _shape_and_supervise(self,
                             tau_direct: np.ndarray,
                             qd_m: np.ndarray,
                             dt: float,
                             bypass_energy_projection: bool = False
                             ) -> tuple[np.ndarray, float]:
        """Apply the strictly ordered haptic output pipeline."""
        tau_saturated, valid = self._shape_candidate(tau_direct, qd_m, dt)
        if not valid:
            tau_zero, diagnostics = self.tank.apply(tau_direct, qd_m, dt)
            self._tau_applied_prev = tau_zero.copy()
            return tau_zero, diagnostics.rho

        # Final energy projection and the only accounting update.
        if bypass_energy_projection:
            tau_app = tau_saturated.copy()
            alpha = 1.0
        else:
            tau_app, diagnostics = self.tank.apply(tau_saturated, qd_m, dt)
            alpha = diagnostics.rho

        # Nothing may modify tau_app after projection.  Only state capture and
        # addition of the independent inverse-dynamics command occur downstream.
        self.commit_applied(tau_app)
        return tau_app, float(alpha)

    def _shape_candidate(
        self,
        tau_direct: np.ndarray,
        qd_m: np.ndarray,
        dt: float,
    ) -> tuple[np.ndarray, bool]:
        """Return the final upstream candidate without applying a supervisor.

        This split is used by the formal PO/PC baseline.  The caller must pass
        the candidate through exactly one supervisor and then invoke
        :meth:`commit_applied`; no additional filtering or limiting is allowed.
        """
        tau_direct = np.asarray(tau_direct, dtype=float).reshape(-1,)
        qd_m = np.asarray(qd_m, dtype=float).reshape(-1,)
        if tau_direct.shape != qd_m.shape:
            raise ValueError("direct haptic torque and qd_m must have equal shapes")

        if self._tau_fb_state.shape != tau_direct.shape:
            self._tau_fb_state = np.zeros_like(tau_direct)
        if self._tau_applied_prev.shape != tau_direct.shape:
            self._tau_applied_prev = np.zeros_like(tau_direct)

        # Invalid samples must not contaminate nominal filter/rate-limit state.
        if (
            not np.isfinite(dt)
            or dt <= 0.0
            or not np.all(np.isfinite(tau_direct))
            or not np.all(np.isfinite(qd_m))
        ):
            self.last_rate_limit_active = False
            self.last_torque_saturation_active = False
            return np.zeros_like(tau_direct), False

        # 1) Optional nominal torque filtering.
        a = float(self.tau_alpha)
        if not 0.0 <= a <= 1.0:
            raise ValueError("tau_alpha must lie in [0, 1]")
        self._tau_fb_state = (1.0 - a) * self._tau_fb_state + a * tau_direct

        # 2) Nominal sign and strength scaling.
        tau_nominal = -self.feedback_strength * self._tau_fb_state

        # 3) Per-joint rate limiting relative to the previously applied output.
        rate_limit = self._symmetric_limit_vector(
            self.torque_rate_limit, tau_nominal.size, "torque_rate_limit"
        )
        max_delta = rate_limit * dt
        tau_rate_limited = np.clip(
            tau_nominal,
            self._tau_applied_prev - max_delta,
            self._tau_applied_prev + max_delta,
        )
        self.last_rate_limit_active = bool(
            np.any(np.abs(tau_rate_limited - tau_nominal) > 1e-12)
        )

        # 4) Per-joint symmetric actuator saturation.
        torque_limit = self._symmetric_limit_vector(
            self.torque_limit, tau_nominal.size, "torque_limit"
        )
        tau_saturated = np.clip(tau_rate_limited, -torque_limit, torque_limit)
        self.last_torque_saturation_active = bool(
            np.any(np.abs(tau_saturated - tau_rate_limited) > 1e-12)
        )
        return tau_saturated, True

    def shape_mapped_reaction_candidate(
        self,
        tau_environment_on_master: np.ndarray,
        qd_m: np.ndarray,
        dt: float,
    ) -> tuple[np.ndarray, bool]:
        """Prepare the common upstream candidate for an external supervisor."""
        reaction = np.asarray(
            tau_environment_on_master,
            dtype=float,
        ).reshape(-1)
        return self._shape_candidate(-reaction, qd_m, dt)

    def commit_applied(self, tau_applied: np.ndarray) -> None:
        """Record the exact supervisor output used at the master port."""
        applied = np.asarray(tau_applied, dtype=float).reshape(-1)
        if applied.shape != self._tau_applied_prev.shape:
            raise ValueError("tau_applied has an unexpected shape")
        if not np.all(np.isfinite(applied)):
            raise ValueError("tau_applied must contain only finite values")
        self._tau_applied_prev = applied.copy()

    # ---------- 高层接口：算 tau_master_current + 反馈 ----------

    def render_mapped_reaction(
        self,
        tau_environment_on_master: np.ndarray,
        qd_m: np.ndarray,
        dt: float,
        *,
        supervise_energy: bool = True,
    ) -> tuple[np.ndarray, float]:
        """Shape an arbitrary generalized reaction such as ``A.T @ tau_s``.

        ``tau_environment_on_master`` already has the desired physical sign:
        it is the environment-on-device reaction and should oppose penetration.
        The historical renderer pipeline stores robot-on-environment action
        internally and inserts its own leading minus sign, so the conversion is
        performed exactly once here.  The returned torque is the final applied
        haptic output; callers must not filter or limit it downstream.
        """
        reaction = np.asarray(
            tau_environment_on_master,
            dtype=float,
        ).reshape(-1,)
        return self._shape_and_supervise(
            -reaction,
            qd_m,
            dt,
            bypass_energy_projection=not supervise_energy,
        )

    def render_tau(self,
                   q_m: np.ndarray,
                   qd_m: np.ndarray,
                   qdd_m: np.ndarray,
                   CF_int_slave_C: np.ndarray,
                   V_slave_C: np.ndarray | None,
                   dt: float,
                   tau_ff_fric: np.ndarray | None = None,
                   ):
        """
        高层接口：给定 (q, qd, qdd) 和从端交互力，返回最终主端总扭矩。

        输入：
        - q_m, qd_m, qdd_m: 主端关节位置/速度/加速度（用于逆动力学）
        - CF_int_slave_C: 从端 InteractionEstimator 估计出的 C F_int（6×1）
        - V_slave_C: 仅为旧调用兼容而保留；能量只按主端
          tau_app^T qd_m 计算
        - tau_ff_fric: 主端摩擦前馈（可为 None）
        - dt: 控制周期（秒）

        输出：
        - tau_cmd_m: 主端最终下发的总扭矩 = tau_master_current + tau_app
        - tau_app: 最终反馈扭矩，也是能量记账使用的精确力矩
        - alpha: 最终能量投影的径向缩放系数
        """
        # 1) 主端逆动力学扭矩
        tau_master_current = self._inverse_dynamics(q_m, qd_m, qdd_m, tau_ff_fric)
        tau_master_current = np.asarray(tau_master_current, dtype=float).reshape(-1,)

        # 2) Render and account the haptic contribution exactly once.
        tau_app, alpha = self.render_from_CF(
            q_m, qd_m,
            CF_int_slave_C=CF_int_slave_C,
            V_slave_C=V_slave_C,
            dt=dt,
        )

        # 3) tau_app is not modified after projection.
        tau_cmd_m = tau_master_current + tau_app
        return tau_cmd_m, tau_app, alpha

    # 兼容旧接口：只关心反馈扭矩
    def render_from_CF(self,
                       q_m: np.ndarray,
                       qd_m: np.ndarray,
                       CF_int_slave_C: np.ndarray,
                       V_slave_C: np.ndarray | None,
                       dt: float):
        """
        兼容旧接口：返回最终施加的反馈扭矩与能量投影系数。
        """
        tau_direct, legacy_alpha = self._compute_feedback_tau(
            q_m,
            qd_m,
            CF_int_slave_C,
            V_slave_C,
            dt,
        )
        return self._shape_and_supervise(
            tau_direct,
            qd_m,
            dt,
            bypass_energy_projection=legacy_alpha is not None,
        )