Blame for predictor.py

#!/usr/bin/env python3
"""Predict real-ball strikes and publish H2 policy-domain commands over ROS 2."""

import time
import math
import argparse
from collections import Counter
import numpy as np
import zmq

import rclpy
from std_msgs.msg import Float64MultiArray


from runtime.prediction_filters import (
    CausalMarkerBallPhaseGate,
    CausalPredictionStabilizer,
    LEGACY_V2_HIT_VELOCITY_BIAS,
)
from runtime.legacy_v2_predictor import HybridTrajectoryPredictor
from runtime.planner_protocol import H2_PLANNER_PACKET_LEN, build_h2_planner_packet
from runtime.provenance import verify_artifacts


def make_latest_sub_socket(context: zmq.Context, endpoint: str) -> zmq.Socket:
    """Create a one-message latest-state SUB socket for combined mocap packets."""
    socket = context.socket(zmq.SUB)
    socket.setsockopt(zmq.CONFLATE, 1)
    socket.setsockopt(zmq.RCVHWM, 1)
    socket.setsockopt(zmq.LINGER, 0)
    socket.setsockopt_string(zmq.SUBSCRIBE, "")
    socket.connect(endpoint)
    return socket


def _source_frames_ready(
    robot_source_monotonic_s: float,
    ball_source_monotonic_s: float,
) -> bool:
    """Return true only after both real mocap sources delivered a first frame."""
    return (
        math.isfinite(float(robot_source_monotonic_s))
        and math.isfinite(float(ball_source_monotonic_s))
        and float(robot_source_monotonic_s) > 0.0
        and float(ball_source_monotonic_s) > 0.0
    )


class PredictionGateDiagnostics:
    """Rate-limited counters for otherwise silent prediction gates."""

    ORDER = (
        "ball_samples",
        "velocity_invalid",
        "phase_waiting",
        "phase_incoming_slow",
        "phase_outgoing",
        "phase_missed",
        "observation_x",
        "observation_y",
        "observation_z",
        "input_pass",
        "predictor_no_hit",
        "predictor_nonfinite",
        "ball_too_low",
        "policy_domain",
        "stabilizer_warmup",
        "stable_prediction",
        "wait_lock_far",
        "lock_too_late",
        "first_lock",
        "lock_refresh",
        "crossed_without_lock",
    )

    def __init__(self, report_interval_s: float) -> None:
        self.report_interval_s = max(0.0, float(report_interval_s))
        self.total = Counter()
        self._last_report = Counter()
        self._last_report_monotonic = time.monotonic()

    def mark(self, reason: str) -> None:
        self.total[reason] += 1

    def report_if_due(self, *, force: bool = False) -> None:
        if self.report_interval_s <= 0.0:
            return
        now = time.monotonic()
        elapsed = now - self._last_report_monotonic
        if not force and elapsed < self.report_interval_s:
            return
        delta = {
            key: self.total[key] - self._last_report[key]
            for key in self.ORDER
            if self.total[key] != self._last_report[key]
        }
        self._last_report = self.total.copy()
        self._last_report_monotonic = now
        if not delta:
            return
        fields = " ".join(f"{key}={value}" for key, value in delta.items())
        print(f"[PREDICT GATES {elapsed:.1f}s] {fields}")

    def report_total(self) -> None:
        if self.report_interval_s <= 0.0 or not self.total:
            return
        fields = " ".join(
            f"{key}={self.total[key]}" for key in self.ORDER if self.total[key]
        )
        print(f"[PREDICT GATES TOTAL] {fields}")


class LegacyV2HitPlanePredictor:
    """Expose the fitted V2 model through this node's hit-plane API."""

    def __init__(self) -> None:
        self._predictor = HybridTrajectoryPredictor()

    def predict_hitting_point(
        self,
        p_init: np.ndarray,
        v_init: np.ndarray,
        target_x: float,
    ):
        trajectory, hit_velocity, flight_time = self._predictor.trajectory_to_x(
            p_init,
            v_init,
            target_x=target_x,
            table_contact_in_range=lambda pos: (
                -1.37 <= float(pos[0]) <= 0.0 and abs(float(pos[1])) <= 0.76
            ),
            max_time=2.0,
            dt=0.005,
        )
        hit_position = trajectory[-1] if hit_velocity is not None else None
        return hit_position, hit_velocity, flight_time


# ==========================================
# 0. Racket command solver in the world frame
# ==========================================
def _solve_v_outgoing(predictor, p_racket: np.ndarray, p_landing: np.ndarray,
                      dt_flight: float, n_iter: int = 4) -> np.ndarray:
    """Newton iteration: find v_out so the ball reaches p_landing in dt_flight
    under the full physics model (drag + spin_az + gravity).

    Starts from the parabolic estimate with g_eff, then corrects via:
        v += (p_target - p_predicted) / dt_flight
    Converges in 3-4 iterations for typical outgoing trajectories (<10 m, <0.6 s).
    """
    g_eff = predictor.g + predictor.spin_az
    v = (p_landing - p_racket) / dt_flight + np.array([0.0, 0.0, 0.5 * g_eff * dt_flight])
    for _ in range(n_iter):
        p_pred, _ = predictor.predict_state_at_time(p_racket, v, dt_flight)
        v = v + (p_landing - p_pred) / dt_flight
    return v


def calculate_racket_command(p_racket: np.ndarray, v_incoming: np.ndarray,
                             p_landing: np.ndarray, dt_flight: float, Cr: float = 0.70,
                             predictor=None):
    """Compute required racket velocity for a given outgoing target.

    A predictor exposing ``g``, ``spin_az`` and ``predict_state_at_time`` uses
    the full iterative flight solve.  The standalone legacy-v2 wrapper only
    predicts trajectories to an x plane, so it deliberately uses the
    parabolic fallback for this compatibility-only racket-velocity field.
    """
    supports_outgoing_solver = (
        predictor is not None
        and hasattr(predictor, "g")
        and hasattr(predictor, "spin_az")
        and callable(getattr(predictor, "predict_state_at_time", None))
    )
    if supports_outgoing_solver:
        v_outgoing = _solve_v_outgoing(predictor, p_racket, p_landing, dt_flight)
    else:
        g_eff = 9.81
        v_outgoing = (p_landing - p_racket) / dt_flight + np.array([0.0, 0.0, 0.5 * g_eff * dt_flight])

    v_diff = v_outgoing - v_incoming
    norm_v_diff = np.linalg.norm(v_diff)

    if norm_v_diff < 1e-6:
        return np.zeros(3)

    u = v_diff / norm_v_diff
    v_racket_mag = (np.dot(v_outgoing, u) + Cr * np.dot(v_incoming, u)) / (1.0 + Cr)
    return v_racket_mag * u


# ==========================================
# 1. Fast ball trajectory predictor with table bounce model
# ==========================================
class FastTrajectoryPredictor:
    def __init__(self, table_height: float = 0.008):
        # Standard 2.7 g ball. Keep Cd/m equal to the 4.0 g calibration so
        # free-flight acceleration remains the trajectory-identified value.
        self.m = 0.0027
        self.g = 9.81
        # Identified from ball.txt using bidirectional, run-level holdout.
        # g + spin_az = 10.150 m/s² for the recorded ball population.
        self.spin_az = 0.340      # topspin Magnus downward accel (m/s²)
        self.Cd_x = 0.000447525
        self.Cd_y = 0.000684450
        self.Cd_z = 0.000590625
        self.table_height = table_height

        # v_rebound = k * v_impact + b, fitted from 503 table impacts.
        self.k = np.array([0.735, 0.446, -0.993])
        self.b = np.array([-0.110, -0.020, 0.0])

    def predict_state_at_time(self, p0: np.ndarray, v0: np.ndarray, t: float):
        x0, y0, z0 = p0
        Vx0, Vy0, Vz0 = v0
        if t < 1e-9: return np.copy(p0), np.copy(v0)

        Vx = self.m * Vx0 / (self.Cd_x * Vx0 * t + self.m) if Vx0 > 0 else self.m * Vx0 / (self.m - self.Cd_x * Vx0 * t)
        x = x0 + (self.m / self.Cd_x) * np.log(1 + self.Cd_x * Vx0 * t / self.m) if Vx0 > 0 else x0 - (self.m / self.Cd_x) * np.log(1 - self.Cd_x * Vx0 * t / self.m)

        Vy = self.m * Vy0 / (self.Cd_y * Vy0 * t + self.m) if Vy0 > 0 else self.m * Vy0 / (self.m - self.Cd_y * Vy0 * t)
        y = y0 + (self.m / self.Cd_y) * np.log(1 + self.Cd_y * Vy0 * t / self.m) if Vy0 > 0 else y0 - (self.m / self.Cd_y) * np.log(1 - self.Cd_y * Vy0 * t / self.m)

        g_eff = self.g + self.spin_az
        mg_Cd = np.sqrt(self.m * g_eff / self.Cd_z)
        sqrt_term = np.sqrt(self.Cd_z * g_eff / self.m)

        if Vz0 >= 0:
            t_peak = (1 / sqrt_term) * np.arctan(Vz0 / mg_Cd) if mg_Cd > 1e-9 else 0.0
            if t <= t_peak:
                Vz = mg_Cd * np.tan(-sqrt_term * t + np.arctan(Vz0 / mg_Cd))
                log_val = np.cos(sqrt_term * t) + (Vz0 / mg_Cd) * np.sin(sqrt_term * t)
                z = z0 + (self.m / self.Cd_z) * np.log(log_val) if log_val > 1e-9 else z0
            else:
                p_peak, v_peak = self.predict_state_at_time(p0, v0, t_peak)
                t_fall = t - t_peak
                exp_term = np.exp(2 * t_fall * sqrt_term)
                Vz = mg_Cd * (1 - exp_term) / (1 + exp_term) + v_peak[2]
                z = p_peak[2] - (self.m / (2 * self.Cd_z)) * np.log(((1 + exp_term) ** 2) / (4 * exp_term)) + v_peak[2] * t_fall
        else:
            ratio_v = np.clip(-Vz0 / mg_Cd, 0.0, 1.0 - 1e-6)
            phase = np.arctanh(ratio_v)
            Vz = -mg_Cd * np.tanh(sqrt_term * t + phase)
            z = z0 + (mg_Cd / sqrt_term) * (np.log(np.cosh(phase)) - np.log(np.cosh(sqrt_term * t + phase)))

        return np.array([x, y, z]), np.array([Vx, Vy, Vz])

    def predict_hitting_point(self, p_init: np.ndarray, v_init: np.ndarray, target_x: float = -1.37, max_time: float = 2.0, dt: float = 0.01):
        t_total = 0.0
        t_seg = 0.0
        max_bounces = 3
        bounce_count = 0

        p0, v0 = np.copy(p_init), np.copy(v_init)
        p_curr, v_curr = p0, v0

        if p0[2] < self.table_height:
            p0[2] = self.table_height
            p_curr = p0.copy()

        while t_total < max_time:
            t_seg += dt
            t_total += dt
            p_next, v_next = self.predict_state_at_time(p0, v0, t_seg)

            if p_curr[0] >= target_x and p_next[0] < target_x:
                ratio = (target_x - p_curr[0]) / (p_next[0] - p_curr[0])
                hit_y = p_curr[1] + ratio * (p_next[1] - p_curr[1])
                hit_z = p_curr[2] + ratio * (p_next[2] - p_curr[2])
                hit_pos = np.array([target_x, hit_y, hit_z])

                hit_vx = v_curr[0] + ratio * (v_next[0] - v_curr[0])
                hit_vy = v_curr[1] + ratio * (v_next[1] - v_curr[1])
                hit_vz = v_curr[2] + ratio * (v_next[2] - v_curr[2])
                hit_vel = np.array([hit_vx, hit_vy, hit_vz])

                hit_t = (t_total - dt) + ratio * dt
                return hit_pos, hit_vel, hit_t

            if p_curr[2] >= self.table_height and p_next[2] < self.table_height:
                ratio = (self.table_height - p_curr[2]) / (p_next[2] - p_curr[2])
                bounce_t_seg = (t_seg - dt) + ratio * dt

                p_impact, v_impact = self.predict_state_at_time(p0, v0, bounce_t_seg)
                p_impact[2] = self.table_height
                v_reb = self.k * v_impact + self.b

                if abs(v_reb[2]) < 0.1:
                    return None, None, None

                bounce_count += 1
                if bounce_count > max_bounces:
                    return None, None, None

                p0, v0 = p_impact, v_reb
                t_seg = 0.0
                t_total = (t_total - dt) + ratio * dt
                p_next, v_next = p0, v0

            p_curr, v_curr = p_next, v_next

        return None, None, None


# ==========================================
# 2. Main loop
# ==========================================

def main(args=None):
    parser = argparse.ArgumentParser(description="Table Tennis Strategy Node")
    parser.add_argument("--sim-ball-only", action="store_true",
                        help="MuJoCo only: ignore robot mocap; publish idle packets before the first ball.")
    parser.add_argument("--topic", default="/wbc_racket_command")
    parser.add_argument("--zmq-endpoint", default="tcp://127.0.0.1:5556",
                        help="Adapter endpoint; one-click deployment passes its private endpoint.")
    parser.add_argument("--debug-log", default=None,
                        help="Optional path for command publish logs, e.g. /tmp/tmp.txt")
    parser.add_argument("--publish-hz", type=float, default=500.0,
                        help="Maximum command publish rate. Use <=0 to disable rate limiting.")
    parser.add_argument(
        "--gate-log-interval",
        type=float,
        default=1.0,
        help="Seconds between prediction-gate summaries. Use <=0 to disable them.",
    )
    parser.add_argument(
        "--robot-base-origin-world", type=float, nargs=3, default=(-1.87, 0.0, -0.76),
        metavar=("X", "Y", "Z"), help="Ready-pose origin in the mocap world frame.",
    )
    parser.add_argument(
        "--hit-plane-x", type=float, default=0.40,
        help="Policy hit plane x in the robot court frame.",
    )
    parser.add_argument(
        "--hit-y-range", type=float, nargs=2, default=(-0.76, 0.76),
        metavar=("MIN", "MAX"), help="Accepted policy hit y range in the robot court frame.",
    )
    parser.add_argument(
        "--hit-z-range", type=float, nargs=2, default=(0.80, 1.50),
        metavar=("MIN", "MAX"), help="Accepted policy hit z range in the robot court frame.",
    )
    parser.add_argument(
        "--lock-time-threshold", type=float, default=0.52,
        help="Start publishing an active strike command below this time-to-hit.",
    )
    parser.add_argument(
        "--prediction-model", choices=("legacy-v2", "deploy-closd"), default="legacy-v2",
        help="Ball flight/bounce model. legacy-v2 matches the default MuJoCo replay model.",
    )
    parser.add_argument(
        "--incoming-velocity-bias", type=float, nargs=3, default=None,
        metavar=("VX", "VY", "VZ"),
        help=(
            "Additive hit-plane velocity calibration in m/s. The default is "
            "the recorded-ball legacy-v2 calibration, or zero for deploy-closd."
        ),
    )
    cli_args, remaining = parser.parse_known_args(args=args)
    hit_y_range = tuple(sorted(map(float, cli_args.hit_y_range)))
    hit_z_range = tuple(sorted(map(float, cli_args.hit_z_range)))
    if hit_y_range[0] >= hit_y_range[1]:
        parser.error("--hit-y-range must have MIN < MAX")
    if hit_z_range[0] >= hit_z_range[1]:
        parser.error("--hit-z-range must have MIN < MAX")
    if cli_args.lock_time_threshold <= 0.0:
        parser.error("--lock-time-threshold must be positive")

    verified = verify_artifacts()
    if "models/real_ball_hybrid_bounce.json" not in verified:
        raise RuntimeError("the default calibrated ball model is not pinned in ARTIFACTS.sha256")

    rclpy.init(args=remaining)
    ros_node = rclpy.create_node('table_tennis_commander_node')

    cmd_pub = ros_node.create_publisher(Float64MultiArray, cli_args.topic, 10)
    ros_node.get_logger().info(
        f"🏓 策略节点启动，将在 {cli_args.topic} 持续发布 H2 v3 指令 (Size=27)..."
    )

    ZMQ_POLL_TIMEOUT_MS = 1

    context = zmq.Context()
    socket = make_latest_sub_socket(context, cli_args.zmq_endpoint)

    poller = zmq.Poller()
    poller.register(socket, zmq.POLLIN)

    print("[系统初始化] 已连接动捕数据源，等待接收数据...")

    predictor = (
        LegacyV2HitPlanePredictor()
        if cli_args.prediction_model == "legacy-v2"
        else FastTrajectoryPredictor(table_height=0.008)
    )
    incoming_velocity_bias = np.asarray(
        cli_args.incoming_velocity_bias
        if cli_args.incoming_velocity_bias is not None
        else (
            LEGACY_V2_HIT_VELOCITY_BIAS
            if cli_args.prediction_model == "legacy-v2"
            else np.zeros(3, dtype=float)
        ),
        dtype=float,
    ).reshape(3)

    ROBOT_BASE_ORIGIN_WORLD = np.asarray(cli_args.robot_base_origin_world, dtype=float)
    HIT_Y_RANGE = hit_y_range
    HIT_Z_RANGE = hit_z_range
    TARGET_HIT_X = float(ROBOT_BASE_ORIGIN_WORLD[0] + cli_args.hit_plane_x)
    LOCK_TIME_THRESHOLD = float(cli_args.lock_time_threshold)
    marker_phase = CausalMarkerBallPhaseGate(target_hit_x=TARGET_HIT_X)
    gate_diagnostics = PredictionGateDiagnostics(cli_args.gate_log_interval)

    p_landing_target = np.array([0.685, 0.0, 0.02])
    desired_flight_time = 0.6

    ball_incoming = False
    has_crossed_target = False
    pending_lock = False

    rally_count = 0
    publish_count = 0

    # Once a valid prediction exists, keep publishing the latest predicted
    # command instead of falling back to default between rallies.
    has_ever_predicted = False

    locked_hit_pos = None
    locked_absolute_hit_time = None
    locked_p_racket_relative = None
    locked_v_racket_cmd_relative = None
    locked_incoming_vel = None

    tentative_hit_pos = None
    tentative_hit_vel = None
    tentative_absolute_hit_time = None
    tentative_stamp = None
    tentative_ball_pos = None
    last_locked_tentative_stamp = None

    prediction_stabilizer = None
    domain_reject_reported = False

    prev_pos_m = None
    prev_mocap_time = None   # Used to interpolate the real crossing time.

    current_robot_base_pos = np.array([-1.8, 0.0, 0.0])
    current_robot_base_quat = np.array([0.0, 0.0, 0.0, 1.0])
    last_ball_t = float("-inf")
    last_robot_t = float("-inf")
    last_robot_source_mono = 0.0
    last_ball_source_mono = 0.0
    last_mocap_time = None
    wall_at_last_mocap = None

    wall_time_at_lock = None
    mocap_time_at_lock = None

    # Default command published until the first valid strike prediction is
    # available.  The current mocap base pose is still forwarded each cycle.
    # Keep the ROS topic in the same post-hit/frozen TTH domain seen by the
    # deployed actor during training.  The active command continues through
    # zero, then saturates at the checkpoint's frozen-command value.
    DEFAULT_T_STRIKE = -0.59
    DEFAULT_P_RACKET_RELATIVE = np.array([cli_args.hit_plane_x, 0.0, 1.0], dtype=float)
    DEFAULT_V_RACKET_RELATIVE = np.array([1.1, 0.0, 0.3], dtype=float)

    print(
        "[策略域] "
        f"model={cli_args.prediction_model}, "
        f"vin_bias={np.array2string(incoming_velocity_bias, precision=3)}, "
        f"world_x={TARGET_HIT_X:.3f}, court_x={cli_args.hit_plane_x:.3f}, "
        f"y=[{HIT_Y_RANGE[0]:.2f},{HIT_Y_RANGE[1]:.2f}], "
        f"z=[{HIT_Z_RANGE[0]:.2f},{HIT_Z_RANGE[1]:.2f}]"
    )
    if cli_args.gate_log_interval > 0.0:
        print(
            f"[PREDICT GATES] 每 {cli_args.gate_log_interval:.1f}s 汇总一次；"
            "计数单位是动捕帧；velocity/phase/observation 表示输入首先未通过的门，"
            "其余是后续预测与锁定状态，不等同于丢球数量。"
        )

    debug_log = open(cli_args.debug_log, "w", buffering=1) if cli_args.debug_log else None
    if debug_log is not None:
        debug_log.write("# timestamp  ball_count  t_strike  "
                        "base_x base_y base_z  qx qy qz qw  "
                        "p_x p_y p_z  v_racket_x v_racket_y v_racket_z  "
                        "v_in_x v_in_y v_in_z\n")

    ros_msg = Float64MultiArray()
    publish_period = 1.0 / cli_args.publish_hz if cli_args.publish_hz > 0.0 else 0.0
    next_publish_mono = 0.0

    try:
        while rclpy.ok():
            latest_robot_msg = None
            latest_ball_msg = None

            if poller.poll(ZMQ_POLL_TIMEOUT_MS):
                try:
                    msg = socket.recv_pyobj(flags=zmq.NOBLOCK)
                except zmq.Again:
                    msg = None

                if msg is not None:
                    if msg.get("type") == "mocap":
                        robot_msg = None if cli_args.sim_ball_only else msg.get("robot")
                        ball_msg = msg.get("ball")

                        if robot_msg is not None:
                            robot_t = float(robot_msg["t"])
                            if robot_t > last_robot_t:
                                latest_robot_msg = robot_msg
                                last_robot_t = robot_t

                        if ball_msg is not None:
                            ball_t = float(ball_msg["t"])
                            if ball_t > last_ball_t:
                                latest_ball_msg = ball_msg
                                last_ball_t = ball_t

            if latest_robot_msg is not None:
                current_robot_base_pos = np.array(latest_robot_msg['pos'], dtype=float) / 1000.0
                current_robot_base_quat = np.array(latest_robot_msg['quat'], dtype=float)

            wall_now = time.time()
            mono_now = time.monotonic()
            if latest_robot_msg is not None:
                last_robot_source_mono = mono_now
            if latest_ball_msg is not None:
                last_ball_source_mono = mono_now

            if latest_ball_msg is not None:
                current_time = latest_ball_msg['t']
                pos_m = np.array(latest_ball_msg['pos'], dtype=float) / 1000.0
                vel_m = np.array(latest_ball_msg['vel'], dtype=float) / 1000.0
                gate_diagnostics.mark("ball_samples")
                phase_transition = marker_phase.update(current_time, pos_m)
                if phase_transition is not None:
                    previous_phase, current_phase = phase_transition
                    print(
                        f"[MARKER PHASE] {previous_phase}->{current_phase} "
                        f"vx={marker_phase.estimated_vx:+.3f}m/s x={pos_m[0]:+.3f}"
                    )

                velocity_valid = bool(latest_ball_msg.get("vel_valid", True))
                prediction_input_ready = False
                if not velocity_valid:
                    gate_diagnostics.mark("velocity_invalid")
                elif not marker_phase.allow_prediction_update:
                    phase_reason = {
                        marker_phase.WAITING: "phase_waiting",
                        marker_phase.INCOMING: "phase_incoming_slow",
                        marker_phase.OUTGOING: "phase_outgoing",
                        marker_phase.MISSED: "phase_missed",
                    }[marker_phase.state]
                    gate_diagnostics.mark(phase_reason)
                elif not TARGET_HIT_X < pos_m[0] < 1.37:
                    gate_diagnostics.mark("observation_x")
                elif not (-1.525 / 2) < pos_m[1] < (1.525 / 2):
                    gate_diagnostics.mark("observation_y")
                elif not 0.0 < pos_m[2] < 1.0:
                    gate_diagnostics.mark("observation_z")
                else:
                    gate_diagnostics.mark("input_pass")
                    prediction_input_ready = True

                if prediction_input_ready:

                    if not ball_incoming:
                        ball_incoming = True
                        has_crossed_target = False
                        pending_lock = False
                        publish_count = -1
                        prev_pos_m = None
                        prev_mocap_time = None
                        tentative_hit_pos = None
                        tentative_hit_vel = None
                        tentative_absolute_hit_time = None
                        tentative_stamp = None
                        tentative_ball_pos = None
                        last_locked_tentative_stamp = None
                        prediction_stabilizer = CausalPredictionStabilizer()
                        domain_reject_reported = False
                        print(f"\n[追踪中] 发现新来球，开始追踪...")

                    if pos_m[0] > TARGET_HIT_X:
                        hit_pos, hit_vel, flight_time = predictor.predict_hitting_point(
                            pos_m, vel_m, target_x=TARGET_HIT_X
                        )
                        if hit_vel is not None:
                            hit_vel = np.asarray(hit_vel, dtype=float) + incoming_velocity_bias

                        if hit_pos is None:
                            gate_diagnostics.mark("predictor_no_hit")
                        elif hit_vel is None or flight_time is None or not (
                            np.isfinite(hit_pos).all()
                            and np.isfinite(hit_vel).all()
                            and math.isfinite(float(flight_time))
                        ):
                            gate_diagnostics.mark("predictor_nonfinite")
                            hit_pos = None
                        elif pos_m[2] < 0.1:
                            gate_diagnostics.mark("ball_too_low")
                            hit_pos = None

                        if hit_pos is not None:
                            hit_pos_relative = hit_pos - ROBOT_BASE_ORIGIN_WORLD
                            in_policy_domain = (
                                HIT_Y_RANGE[0] <= hit_pos_relative[1] <= HIT_Y_RANGE[1]
                                and HIT_Z_RANGE[0] <= hit_pos_relative[2] <= HIT_Z_RANGE[1]
                            )
                            if not in_policy_domain:
                                gate_diagnostics.mark("policy_domain")
                                if not domain_reject_reported:
                                    print(
                                        "[PREDICT REJECT] 击球点超出策略域 | "
                                        f"y={hit_pos_relative[1]:+.3f} "
                                        f"z={hit_pos_relative[2]:+.3f}"
                                    )
                                    domain_reject_reported = True
                                hit_pos = None

                        if hit_pos is not None and prediction_stabilizer is not None:
                            # Use exactly the same causal median/innovation/EMA
                            # filter as recorded-ball MuJoCo replay.  Filtering
                            # absolute impact time preserves a smooth countdown;
                            # filtering hit velocity removes the remaining command
                            # jerk that the former position-only EMA did not cover.
                            stable = prediction_stabilizer.update(
                                current_time,
                                current_time + flight_time,
                                hit_pos,
                                hit_vel,
                            )
                            if stable is None:
                                gate_diagnostics.mark("stabilizer_warmup")
                                gate_diagnostics.report_if_due()
                                continue
                            gate_diagnostics.mark("stable_prediction")
                            (
                                tentative_absolute_hit_time,
                                tentative_hit_pos,
                                tentative_hit_vel,
                            ) = stable
                            tentative_stamp = current_time
                            tentative_ball_pos = pos_m.copy()

                # Post-check when the real ball crosses the target_x plane.
                if ball_incoming and not has_crossed_target and prev_pos_m is not None:
                    if prev_pos_m[0] >= TARGET_HIT_X and pos_m[0] < TARGET_HIT_X:
                        has_crossed_target = True
                        ball_incoming = False

                        # Interpolate the real crossing position.
                        ratio = (TARGET_HIT_X - prev_pos_m[0]) / (pos_m[0] - prev_pos_m[0])
                        real_hit_y = prev_pos_m[1] + ratio * (pos_m[1] - prev_pos_m[1])
                        real_hit_z = prev_pos_m[2] + ratio * (pos_m[2] - prev_pos_m[2])

                        # Interpolate the real crossing mocap timestamp.
                        if prev_mocap_time is not None:
                            real_hit_time = prev_mocap_time + ratio * (current_time - prev_mocap_time)
                        else:
                            real_hit_time = current_time  # fallback

                        print(f"[Ball #{rally_count}] 事后校验 | "
                              f"真实穿越 Y={real_hit_y:.4f} Z={real_hit_z:.4f} "
                              f"T={real_hit_time:.4f}s", end="")
                        if pending_lock and locked_hit_pos is not None and locked_absolute_hit_time is not None:
                            error_y = real_hit_y - locked_hit_pos[1]
                            error_z = real_hit_z - locked_hit_pos[2]
                            time_error_ms = (locked_absolute_hit_time - real_hit_time) * 1000
                            print(f" | 位置误差 {math.hypot(error_y, error_z) * 1000:.1f}mm"
                                  f" (dY={error_y * 1000:+.1f} dZ={error_z * 1000:+.1f})"
                                  f" | 时间误差 {time_error_ms:+.1f}ms")
                        else:
                            gate_diagnostics.mark("crossed_without_lock")
                            print(" | [PREDICT MISS] 穿越击球面前未锁定命令")

                if marker_phase.terminal and ball_incoming:
                    ball_incoming = False

                prev_pos_m = pos_m.copy()
                prev_mocap_time = current_time   # Previous mocap timestamp.
                last_mocap_time = current_time
                wall_at_last_mocap = wall_now

            if ball_incoming and tentative_absolute_hit_time is not None \
                    and last_mocap_time is not None and wall_at_last_mocap is not None:

                if latest_ball_msg is not None:
                    estimated_mocap_now = last_mocap_time
                else:
                    estimated_mocap_now = last_mocap_time + (wall_now - wall_at_last_mocap)

                t_to_hit = tentative_absolute_hit_time - estimated_mocap_now
                first_lock = (not pending_lock) and (0.0 < t_to_hit <= LOCK_TIME_THRESHOLD)
                refresh_lock = pending_lock and (tentative_stamp != last_locked_tentative_stamp)
                if latest_ball_msg is not None and not pending_lock:
                    if t_to_hit > LOCK_TIME_THRESHOLD:
                        gate_diagnostics.mark("wait_lock_far")
                    elif t_to_hit <= 0.0:
                        gate_diagnostics.mark("lock_too_late")

                if first_lock or refresh_lock:
                    locked_hit_pos = tentative_hit_pos.copy()
                    locked_absolute_hit_time = tentative_absolute_hit_time
                    locked_incoming_vel = tentative_hit_vel.copy()

                    locked_v_racket_cmd = calculate_racket_command(
                        p_racket=locked_hit_pos,
                        v_incoming=tentative_hit_vel,
                        p_landing=p_landing_target,
                        dt_flight=desired_flight_time,
                        Cr=0.70,
                        predictor=predictor,
                    )

                    locked_p_racket_relative = locked_hit_pos - ROBOT_BASE_ORIGIN_WORLD
                    locked_v_racket_cmd_relative = locked_v_racket_cmd.copy()
                    # locked_v_racket_cmd_relative[0] += 0.6

                    has_ever_predicted = True
                    last_locked_tentative_stamp = tentative_stamp
                    wall_time_at_lock = wall_now
                    mocap_time_at_lock = estimated_mocap_now

                    if first_lock:
                        gate_diagnostics.mark("first_lock")
                        pending_lock = True
                        publish_count = 0
                        gap_note = "" if latest_ball_msg is not None else "（丢帧期间由时间推算触发）"
                        bp = tentative_ball_pos
                        print(f"[追踪中] 首次锁定{gap_note} | t_to_hit={t_to_hit:.3f}s | "
                              f"击球点 Y={locked_hit_pos[1]:.3f} Z={locked_hit_pos[2]:.3f} | "
                              f"球位 ({bp[0]:+.3f}, {bp[1]:+.3f}, {bp[2]:+.3f})")
                    else:
                        gate_diagnostics.mark("lock_refresh")

            gate_diagnostics.report_if_due()
            # The planner-v3 packet deliberately rejects absent/non-positive
            # source timestamps.  Keep the zero sentinels until real frames
            # arrive and publish nothing until *both* sources have done so;
            # substituting the loop clock here would fabricate freshness.
            if not cli_args.sim_ball_only and not _source_frames_ready(last_robot_source_mono, last_ball_source_mono):
                continue
            if publish_period > 0.0 and mono_now < next_publish_mono:
                continue
            next_publish_mono = mono_now + publish_period

            # Preserve the legacy first 14 values and append the fields needed
            # by the learned policy: ball count, predicted incoming velocity at
            # the hit plane, and the landing target in the robot court frame.
            # Before the first prediction, publish default pose/velocity with
            # t_strike=-0.59.  Afterward, keep using the latest locked command
            # and let t_strike decay down to DEFAULT_T_STRIKE.
            publish_t_strike = DEFAULT_T_STRIKE
            p_cmd_publish = DEFAULT_P_RACKET_RELATIVE
            v_cmd_publish = DEFAULT_V_RACKET_RELATIVE
            incoming_vel_publish = np.zeros(3, dtype=float)
            has_prediction = (
                has_ever_predicted
                and locked_hit_pos is not None
                and locked_absolute_hit_time is not None
            )

            if has_prediction:
                if latest_ball_msg is not None:
                    dynamic_t_strike = locked_absolute_hit_time - current_time
                else:
                    elapsed_wall = wall_now - wall_time_at_lock
                    estimated_mocap_time = mocap_time_at_lock + elapsed_wall
                    dynamic_t_strike = locked_absolute_hit_time - estimated_mocap_time

                # Clamp outward t_strike while keeping the latest predicted
                # command after the strike time has passed.
                publish_t_strike = float(np.clip(dynamic_t_strike, DEFAULT_T_STRIKE, 0.8))

                p_cmd_publish = np.array([
                    float(cli_args.hit_plane_x),
                    float(locked_p_racket_relative[1]),
                    float(locked_p_racket_relative[2]),
                ])
                # Work on a copy: this block runs once per publish cycle, so
                # incrementing the stored command in place accumulated the offset
                # a non-deterministic number of times between mocap frames.
                v_cmd_raw = np.array(locked_v_racket_cmd_relative, dtype=float)
                incoming_vel_publish = np.array(locked_incoming_vel, dtype=float)
                if abs(p_cmd_publish[1]) > 0.20 and p_cmd_publish[2] < 0.91:
                    v_cmd_raw[0] += 0.4
                if p_cmd_publish[2] < 0.91:
                    v_cmd_publish = np.array([
                        np.clip(v_cmd_raw[0], 1.31, 2.99),
                        np.clip(v_cmd_raw[1], -0.09, 0.09),
                        np.clip(v_cmd_raw[2], 0.01, 1.09)
                    ])
                else:
                    v_cmd_publish = np.array([
                        np.clip(v_cmd_raw[0], 1.31, 2.99),
                        np.clip(v_cmd_raw[1], -0.09, 0.09),
                        np.clip(v_cmd_raw[2], 0.01, 0.39)
                    ])

            published_ball_count = rally_count + int(has_prediction and publish_count == 0)
            landing_xy_relative = p_landing_target[:2] - ROBOT_BASE_ORIGIN_WORLD[:2]
            cur_msg = build_h2_planner_packet(
                robot_pos_world=current_robot_base_pos,
                robot_quat_xyzw=current_robot_base_quat,
                time_to_hit_s=publish_t_strike,
                hit_pos_robot=p_cmd_publish,
                racket_velocity_robot=v_cmd_publish,
                command_id=published_ball_count,
                incoming_velocity_robot=incoming_vel_publish,
                landing_xy_robot=landing_xy_relative,
                robot_source_monotonic_s=last_robot_source_mono,
                ball_source_monotonic_s=last_ball_source_mono,
                robot_base_origin_world=ROBOT_BASE_ORIGIN_WORLD,
                allow_missing_sources=cli_args.sim_ball_only,
            ).tolist()

            if len(cur_msg) != H2_PLANNER_PACKET_LEN:
                raise RuntimeError(f"internal H2 planner packet length changed: {len(cur_msg)}")

            ros_msg.data = cur_msg
            cmd_pub.publish(ros_msg)

            if debug_log is not None:
                debug_log.write(
                    f"{wall_now:.4f}  ball={published_ball_count}  "
                    f"t_strike={publish_t_strike:.4f}  "
                    f"base=[{cur_msg[0]:.3f},{cur_msg[1]:.3f},{cur_msg[2]:.3f}]  "
                    f"quat=[{cur_msg[3]:.3f},{cur_msg[4]:.3f},{cur_msg[5]:.3f},{cur_msg[6]:.3f}]  "
                    f"p=[{p_cmd_publish[0]:.3f},{p_cmd_publish[1]:.3f},{p_cmd_publish[2]:.3f}]  "
                    f"v_racket=[{v_cmd_publish[0]:.3f},{v_cmd_publish[1]:.3f},{v_cmd_publish[2]:.3f}]  "
                    f"v_in=[{incoming_vel_publish[0]:.3f},{incoming_vel_publish[1]:.3f},{incoming_vel_publish[2]:.3f}]\n"
                )

            if has_prediction and publish_count >= 0:
                publish_count += 1
                if publish_count == 1:
                    rally_count += 1
                    print(f"\n[Ball #{rally_count}] 首次发布成功 | "
                          f"击球点 Y={locked_p_racket_relative[1]:.3f} "
                          f"Z={locked_p_racket_relative[2]:.3f}")

    except KeyboardInterrupt:
        print("\n\n[系统退出] 检测到中止信号 (Ctrl+C)，正在安全关闭。")
    finally:
        gate_diagnostics.report_total()
        if debug_log is not None:
            debug_log.close()
        socket.close(0)
        context.term()
        ros_node.destroy_node()
        rclpy.try_shutdown()


if __name__ == "__main__":
    main()