Add Wall-X serving and Turtle2 TCP WebSocket bridge

tests/test_lora_scale.py

@@ -0,0 +1,45 @@
+"""Keep serving LoRA merge consistent with exported training parameters."""
+
+import json
+
+import pytest
+import torch
+
+from wall_x._vendor.harrix.utils import ckpt_load
+
+
+def _tiny_lora():
+    stem = "model.base_model.model.proj"
+    return {
+        stem + ".base_layer.weight": torch.zeros((2, 2)),
+        stem + ".lora_A.default.weight": torch.eye(2),
+        stem + ".lora_B.default.weight": torch.eye(2),
+    }
+
+
+def test_checkpoint_local_lora_config_controls_merge(tmp_path):
+    (tmp_path / "lora_config.json").write_text(
+        json.dumps({"lora_r": 2, "lora_alpha": 6}), encoding="utf-8"
+    )
+    tensors = _tiny_lora()
+    scale = ckpt_load.resolve_lora_scale(str(tmp_path), {}, tensors)
+    merged = ckpt_load.reshape_compatible_state_dict(
+        tensors, {"model.proj.weight": torch.zeros((2, 2))}, lora_scale=scale
+    )
+    assert scale == 3
+    torch.testing.assert_close(merged["model.proj.weight"], 3 * torch.eye(2))
+
+
+def test_checkpoint_lora_rank_mismatch_is_rejected(tmp_path):
+    (tmp_path / "lora_config.json").write_text(
+        json.dumps({"lora_r": 4, "lora_alpha": 8}), encoding="utf-8"
+    )
+    with pytest.raises(ValueError, match="rank"):
+        ckpt_load.resolve_lora_scale(str(tmp_path), {}, _tiny_lora())
+
+
+def test_legacy_checkpoint_without_lora_metadata_uses_previous_scale(tmp_path):
+    messages = []
+    scale = ckpt_load.resolve_lora_scale(str(tmp_path), {}, _tiny_lora(), log_fn=messages.append)
+    assert scale == 2
+    assert any("metadata" in message for message in messages)

tests/test_robot_preprocessor.py

@@ -0,0 +1,29 @@
+"""Online robot preprocessing checks shared by Turtle serving."""
+
+from types import SimpleNamespace
+
+import numpy as np
+
+from wall_x._vendor.harrix.serving._wallx_infer.robot import Robot
+
+
+def test_dof_mask_disables_virtual_action_padding():
+    wrapper = SimpleNamespace(
+        config=SimpleNamespace(
+            action_horizon=10,
+            train_config={
+                "dof_config": {
+                    "master_right_ee_cartesian_pos": 3,
+                    "master_right_ee_rotation": 3,
+                    "master_right_gripper": 1,
+                    "action_padding": 19,
+                }
+            },
+        )
+    )
+
+    mask = Robot._get_dof_mask(wrapper)
+
+    assert mask.shape == (1, 10, 26)
+    np.testing.assert_array_equal(mask[:, :, :7], 1)
+    np.testing.assert_array_equal(mask[:, :, 7:], 0)

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])
+

tests/test_run_serving_contract.py

@@ -0,0 +1,40 @@
+"""Check that the launch command rejects a mismatched model action length."""
+
+import os
+from pathlib import Path
+import subprocess
+import pytest
+
+
+ROOT = Path(__file__).parents[1]
+SCRIPT = ROOT / "scripts" / "run_serving.sh"
+
+
+def _launch(horizon):
+    if "WALLX_TEST_CHECKPOINT" not in os.environ:
+        pytest.skip("set WALLX_TEST_CHECKPOINT to run checkpoint-specific contract tests")
+    checkpoint = Path(os.environ["WALLX_TEST_CHECKPOINT"])
+    env = os.environ.copy()
+    env["PYTHON_BIN"] = os.environ.get("WALLX_TEST_PYTHON", "python")
+    return subprocess.run(
+        [
+            "bash", str(SCRIPT),
+            "--checkpoint-path", str(checkpoint),
+            "--train-config-path", str(checkpoint / "config.yml"),
+            "--action-horizon", str(horizon),
+            "--dry-run",
+        ],
+        cwd=ROOT, env=env, text=True, capture_output=True, check=False,
+    )
+
+
+def test_launch_rejects_mismatch_to_training_horizon():
+    result = _launch(10)
+    assert result.returncode != 0
+    assert "action horizon" in (result.stderr + result.stdout).lower()
+
+
+def test_launch_accepts_training_horizon():
+    result = _launch(32)
+    assert result.returncode == 0, result.stderr
+    assert "--model-config.action-horizon  32" in result.stdout

tests/test_tcp_ws_bridge.py

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

tests/test_turtle2_inference_input.py

@@ -0,0 +1,47 @@
+"""Checkpoint-specific multimodal prompt contract, without loading model weights."""
+
+import logging
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from wall_x._vendor.harrix.serving._wallx_infer.infer_config import InferConfig
+from wall_x._vendor.harrix.serving._wallx_infer.model_wrapper import WallxModelWrapper
+from wall_x.trainer.trainer_utils import load_wallx_processors
+
+
+def test_three_square_cameras_preserve_all_image_tokens():
+    """Truncating an image placeholder must never reach the VLA as valid input."""
+    if "WALLX_TEST_CHECKPOINT" not in os.environ:
+        pytest.skip("set WALLX_TEST_CHECKPOINT to run checkpoint-specific input tests")
+    checkpoint = Path(os.environ["WALLX_TEST_CHECKPOINT"])
+    config = InferConfig(
+        checkpoint_path=str(checkpoint),
+        train_config_path=str(checkpoint / "config.yml"),
+        model_device="cpu", action_horizon=32,
+    )
+    processor = load_wallx_processors(config.train_config, device="cpu")["processor"]
+    wrapper = WallxModelWrapper.__new__(WallxModelWrapper)
+    wrapper.config = config
+    wrapper.model = SimpleNamespace(processor=processor)
+    wrapper.cam_names = config.cam_names
+    wrapper.norm_key = "pick_paper_lerobot_v21_3cam"
+    wrapper.logger = logging.getLogger(__name__)
+    wrapper.role_start_symbol = "<|im_start|>"
+    wrapper.role_end_symbol = "<|im_end|>"
+    wrapper.vision_start_symbol = "<|vision_start|>"
+    wrapper.vision_end_symbol = "<|vision_end|>"
+    wrapper.image_pad_symbol = "<|image_pad|>"
+    wrapper.propri_symbol = "<|propri|>"
+    wrapper.action_symbol = "<|action|>"
+    prefix, postfix = wrapper.get_text_for_action("pick up the paper towel")
+    image = np.full((448, 448, 3), 127, dtype=np.uint8)
+    observation = [{name: image for name in wrapper.cam_names}]
+    inputs = wrapper.construct_model_input(observation, prefix, postfix)
+    merge = processor.image_processor.merge_size ** 2
+    expected = int((inputs["image_grid_thw"].prod(dim=1) // merge).sum())
+    actual = int((inputs["input_ids"] == processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")).sum())
+    assert actual == expected == 768