VLA/tests/test_rtc_wallx.py

"""Offline tests for the Wall-X RTC copies."""

import importlib.util
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import torch

from wall_x._vendor.harrix.serving.rtc_wallx import (
    WallXRTCConfig,
    WallXRTCProcessor,
)
from wall_x._vendor.harrix.serving.policy.wall_x_policy_rtc import WallXPolicy
from wall_x._vendor.x2robot_utils import geometry as geom


ROOT = Path(__file__).parents[1]


def _load_rtc_bridge():
    path = ROOT / "scripts" / "tcp_ws_bridge_rtc.py"
    spec = importlib.util.spec_from_file_location("tcp_ws_bridge_rtc", path)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


def test_linear_prefix_weights():
    processor = WallXRTCProcessor(
        WallXRTCConfig(execution_horizon=6, prefix_attention_schedule="linear")
    )
    weights = processor.get_prefix_weights(start=2, end=6, total=10)

    torch.testing.assert_close(weights[:2], torch.ones(2))
    assert torch.all(weights[2:6] < 1)
    assert torch.all(weights[2:6] > 0)
    torch.testing.assert_close(weights[6:], torch.zeros(4))


def test_no_prefix_preserves_wallx_velocity():
    processor = WallXRTCProcessor(WallXRTCConfig())
    x = torch.randn(1, 8, 4)

    result = processor.guide_increasing_flow(
        x_t=x,
        time=torch.tensor(0.4),
        predict_velocity=lambda value: torch.ones_like(value),
        prev_chunk_left_over=None,
    )

    torch.testing.assert_close(result, torch.ones_like(x))


def test_guidance_changes_velocity_toward_prefix():
    processor = WallXRTCProcessor(
        WallXRTCConfig(
            execution_horizon=6,
            max_guidance_weight=10.0,
            prefix_attention_schedule="ones",
        )
    )
    x = torch.zeros(1, 8, 3)
    prefix = torch.ones(1, 8, 3)

    result = processor.guide_increasing_flow(
        x_t=x,
        time=torch.tensor(0.5),
        predict_velocity=lambda value: value * 0,
        prev_chunk_left_over=prefix,
        inference_delay=0,
        execution_horizon=6,
    )

    assert torch.all(result[:, :6] > 0)
    torch.testing.assert_close(result[:, 6:], torch.zeros_like(result[:, 6:]))


def test_rtc_bridge_keeps_full_horizon_before_interpolation():
    bridge = _load_rtc_bridge()
    horizon = 32
    multiplier = 6
    right = np.zeros((horizon, 7), dtype=np.float64)
    right[:, 0] = np.arange(horizon, dtype=np.float64)

    actions = bridge.prepare_robot_actions(
        {"follow2_pos": right.tolist()},
        state_follow1_pos=np.zeros(7),
        state_follow2_pos=np.zeros(7),
        state_head_pos=[0.0, -1.0],
        state_lift=[0.3],
        state_car_pose=np.zeros(3),
        action_horizon=horizon,
        action_end_ratio=1.0,
        action_interpolate_multiplier=multiplier,
        max_position_delta=100.0,
        max_rotation_delta=100.0,
    )

    assert len(actions["follow2_pos"]) == horizon * multiplier
    assert actions["follow2_pos"][-1][0] == 31.0

def test_relative_prefix_is_reanchored_in_normalized_model_layout():
    class IdentityNormalizer:
        @staticmethod
        def normalize_data(value, dataset_names):
            assert dataset_names == ['test']
            return value

    policy = WallXPolicy.__new__(WallXPolicy)
    policy.config = SimpleNamespace(
        model_device='cpu',
        train_config={
            'dof_config': {
                'action_padding': 16,
                'follow_right_ee_cartesian_pos_relative': 3,
                'follow_right_ee_rotation_6D_relative': 6,
                'follow_right_gripper': 1,
            }
        },
    )
    policy.model_wrapper = SimpleNamespace(
        normalizer_action=IdentityNormalizer(), norm_key='test'
    )
    current = np.array([0.4, -0.2, 0.3, 0.1, -0.2, 0.3, 0.5])
    target = np.array([0.45, -0.1, 0.28, 0.1, -0.2, 0.3, 0.7])
    policy._rtc_sessions = {
        'robot': {'follow2_pos': np.stack([current, target]), 'chunk_id': 1}
    }

    prefix = policy._rtc_build_normalized_prefix(
        {'follow2_pos': current.tolist()},
        {'session_id': 'robot', 'consumed_model_steps': 1},
    )

    assert prefix.shape == (1, 1, 26)
    torch.testing.assert_close(prefix[0, 0, :16], torch.zeros(16))
    torch.testing.assert_close(
        prefix[0, 0, 16:19],
        torch.tensor(target[:3] - current[:3], dtype=torch.float32),
    )
    expected_identity_6d = torch.tensor(
        geom.euler_to_matrix_zyx_6d_nb(np.zeros((1, 3)))[0],
        dtype=torch.float32,
    )
    torch.testing.assert_close(prefix[0, 0, 19:25], expected_identity_6d)
    assert np.isclose(prefix[0, 0, 25].item(), target[6])