VLA/tests/test_tcp_ws_bridge.py

"""Safety and protocol checks for the Turtle2 action bridge."""

import importlib.util
from pathlib import Path

import numpy as np
import pytest


_BRIDGE_PATH = Path(__file__).parents[1] / "scripts" / "tcp_ws_bridge.py"
_SPEC = importlib.util.spec_from_file_location("tcp_ws_bridge", _BRIDGE_PATH)
bridge = importlib.util.module_from_spec(_SPEC)
assert _SPEC.loader is not None
_SPEC.loader.exec_module(bridge)




def test_packet_limit_scales_chunk_without_cumulative_ramp():
    current = np.zeros(7)
    traj = np.array(
        [
            [0.10, 0.04, -0.02, 0.20, -0.10, 0.05, 0.0],
            [0.20, 0.08, -0.04, 0.40, -0.20, 0.10, 0.1],
            [0.40, 0.16, -0.08, 0.80, -0.40, 0.20, 0.2],
        ]
    )

    limited, scales = bridge._limit_action_packet(traj, current, 0.005, 0.0125)

    np.testing.assert_allclose(limited[:, :3], traj[:, :3] * 0.005 / 0.40)
    np.testing.assert_allclose(limited[:, 3:6], traj[:, 3:6] * 0.0125 / 0.80)
    assert np.max(np.abs(limited[:, :3] - current[:3])) <= 0.005 + 1e-12
    assert np.max(np.abs(limited[:, 3:6] - current[3:6])) <= 0.0125 + 1e-12
    assert np.max(np.abs(np.diff(limited[:, :3], axis=0))) <= 0.005 + 1e-12
    assert np.max(np.abs(np.diff(limited[:, 3:6], axis=0))) <= 0.0125 + 1e-12
    assert scales["position"][1] < 1.0
    assert scales["rotation"][1] < 1.0


def test_prepare_robot_actions_keeps_packet_endpoint_small():
    traj = np.array(
        [
            [0.10, 0.00, 0.00, 0.20, 0.00, 0.00, 0.0],
            [0.20, 0.00, 0.00, 0.40, 0.00, 0.00, 0.0],
            [0.30, 0.00, 0.00, 0.60, 0.00, 0.00, 0.0],
            [0.40, 0.00, 0.00, 0.80, 0.00, 0.00, 0.0],
        ]
    )
    state = np.zeros(7)
    actions = bridge.prepare_robot_actions(
        {"follow1_pos": traj.tolist(), "follow2_pos": traj.tolist()},
        state_follow1_pos=state,
        state_follow2_pos=state,
        state_head_pos=[0.0, -1.0],
        state_lift=[0.4],
        state_car_pose=np.zeros(3),
        action_horizon=4,
        action_end_ratio=1.0,
        action_interpolate_multiplier=1,
        max_position_delta=0.005,
        max_rotation_delta=0.0125,
        clip_action_delta=True,
    )

    right = np.asarray(actions["follow2_pos"])
    assert right[-1, 0] <= 0.005 + 1e-12
    assert right[-1, 3] <= 0.0125 + 1e-12
    assert right[-1, 0] < 0.01  # old cumulative clipping reached 4 * 0.005


def test_disabled_base_motion_sends_relative_zero():
    actions = bridge.prepare_robot_actions(
        {"follow1_pos": [[0.0] * 7] * 2, "follow2_pos": [[0.0] * 7] * 2},
        state_follow1_pos=[0.0] * 7, state_follow2_pos=[0.0] * 7,
        state_head_pos=[0.0, -1.0], state_lift=[0.1], state_car_pose=[1.2, -0.3, 0.7],
        action_horizon=2, action_end_ratio=1.0, action_interpolate_multiplier=1,
        max_position_delta=0.005, max_rotation_delta=0.0125, clip_action_delta=True,
    )
    assert actions["car_pose"] == [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]


def test_right_arm_only_mode_holds_serializer_left_arm():
    state_left = np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.4])
    right = np.array(
        [
            [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.7],
            [0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.8],
        ]
    )
    synthetic_left = np.full((2, 7), 99.0)
    actions = bridge.prepare_robot_actions(
        {"follow1_pos": synthetic_left.tolist(), "follow2_pos": right.tolist()},
        state_follow1_pos=state_left,
        state_follow2_pos=np.zeros(7),
        state_head_pos=[0.0, -1.0],
        state_lift=[0.4],
        state_car_pose=np.zeros(3),
        action_horizon=2,
        action_end_ratio=1.0,
        action_interpolate_multiplier=1,
    )
    np.testing.assert_allclose(actions["follow1_pos"], np.repeat(state_left[None, :], 2, axis=0))
    np.testing.assert_allclose(actions["follow2_pos"], right)


def test_right_arm_feedback_reports_missing_motion():
    initial = np.zeros(7)
    command = np.array([0.003, 0.0, 0.0, 0.0, 0.01, 0.0, 0.5])
    missing = bridge.assess_right_arm_feedback(initial, initial, command)
    assert missing["command_requests_motion"]
    assert missing["missing_feedback"]

    observed = np.array([0.001, 0.0, 0.0, 0.0, 0.003, 0.0, 0.1])
    moved = bridge.assess_right_arm_feedback(initial, observed, command)
    assert moved["observed_motion"]
    assert not moved["missing_feedback"]


def test_raw_action_response_is_rejected_before_robot_packet():
    """The unsupported raw path must not reinterpret virtual padding as pose."""
    raw = np.zeros((2, 26))
    raw[:, :10] = [0.01, 0.02, 0.03, 1, 0, 0, 0, 1, 0, 0.8]
    raw[:, 16:26] = [0.20, 0.10, 0.05, 1, 0, 0, 0, 1, 0, 0.2]

    with pytest.raises(ValueError, match="--serialize-actions"):
        bridge.prepare_robot_actions(
            {"predict_action": raw.tolist()},
            state_follow1_pos=np.zeros(7),
            state_follow2_pos=np.zeros(7),
            state_head_pos=[0, -1],
            state_lift=[0.1],
            state_car_pose=np.zeros(3),
            action_horizon=2,
            action_end_ratio=1,
            action_interpolate_multiplier=1,
        )