Add Wall-X serving and Turtle2 TCP WebSocket bridge

tests/test_rtc_wallx.py

@@ -0,0 +1,149 @@
+"""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])
+