150 lines
4.6 KiB
Python
150 lines
4.6 KiB
Python
"""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])
|
|
|