File size: 5,305 Bytes
31ade1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | from __future__ import annotations
import numpy as np
from eval.dual_push_retarget_utils import (
backproject_depth_to_world,
delta_action_from_absolute_targets,
estimate_dual_push_target_positions_from_front_rgbd,
extract_dual_push_button_positions,
make_bimanual_absolute_action,
parse_dual_push_target_colors,
retarget_demo_pose_to_live_button,
)
def test_extract_dual_push_button_positions_reads_expected_slices() -> None:
state = np.arange(136, dtype=np.float32)
button0, button1, button2 = extract_dual_push_button_positions(state)
assert button0.tolist() == [7.0, 8.0, 9.0]
assert button1.tolist() == [14.0, 15.0, 16.0]
assert button2.tolist() == [21.0, 22.0, 23.0]
def test_retarget_demo_pose_to_live_button_translates_position_only() -> None:
demo_pose = np.array([0.4, 0.5, 0.6, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
demo_button = np.array([0.1, 0.2, 0.3], dtype=np.float32)
live_button = np.array([0.3, -0.1, 0.7], dtype=np.float32)
retargeted = retarget_demo_pose_to_live_button(demo_pose, demo_button, live_button)
np.testing.assert_allclose(retargeted[:3], np.array([0.6, 0.2, 1.0], dtype=np.float32))
np.testing.assert_allclose(retargeted[3:], demo_pose[3:])
def test_make_bimanual_absolute_action_packs_expected_layout() -> None:
right_pose = np.array([1, 2, 3, 0, 0, 0, 1], dtype=np.float32)
left_pose = np.array([4, 5, 6, 0, 0, 1, 0], dtype=np.float32)
action = make_bimanual_absolute_action(right_pose, left_pose, 1.0, 0.0)
assert action.shape == (18,)
np.testing.assert_allclose(action[:7], right_pose)
np.testing.assert_allclose(action[9:16], left_pose)
assert action[7] == 1.0
assert action[8] == 1.0
assert action[16] == 0.0
assert action[17] == 1.0
def test_delta_action_from_absolute_targets_matches_translation_and_gripper() -> None:
current_right = np.array([0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
current_left = np.array([-0.1, 0.4, 0.2, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
target_right = np.array([0.15, 0.1, 0.35, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
target_left = np.array([-0.05, 0.5, 0.25, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
delta = delta_action_from_absolute_targets(
current_right,
current_left,
target_right,
target_left,
right_gripper_open=1.0,
left_gripper_open=0.0,
)
np.testing.assert_allclose(delta[:3], np.array([0.05, -0.1, 0.05], dtype=np.float32), atol=1e-6)
np.testing.assert_allclose(delta[7:10], np.array([0.05, 0.1, 0.05], dtype=np.float32), atol=1e-6)
np.testing.assert_allclose(delta[3:6], np.zeros(3, dtype=np.float32), atol=1e-6)
np.testing.assert_allclose(delta[10:13], np.zeros(3, dtype=np.float32), atol=1e-6)
assert delta[6] == 1.0
assert delta[13] == 0.0
def test_delta_action_from_absolute_targets_captures_rotation_delta() -> None:
current_right = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
current_left = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0], dtype=np.float32)
half_turn_z_xyzw = np.array([0.0, 0.0, np.sin(np.pi / 4.0), np.cos(np.pi / 4.0)], dtype=np.float32)
target_right = np.concatenate([np.zeros(3, dtype=np.float32), half_turn_z_xyzw], axis=0)
target_left = current_left.copy()
delta = delta_action_from_absolute_targets(
current_right,
current_left,
target_right,
target_left,
right_gripper_open=1.0,
left_gripper_open=1.0,
)
np.testing.assert_allclose(delta[:3], np.zeros(3, dtype=np.float32), atol=1e-6)
np.testing.assert_allclose(delta[7:13], np.zeros(6, dtype=np.float32), atol=1e-6)
assert abs(float(delta[5])) > 1.4
def test_parse_dual_push_target_colors_reads_both_targets() -> None:
assert parse_dual_push_target_colors("push the olive and the orange buttons") == ("olive", "orange")
def test_backproject_depth_to_world_handles_signed_focal_lengths() -> None:
depth = np.array([[0.75]], dtype=np.float32)
intrinsics = np.array([[-100.0, 0.0, 0.0], [0.0, -100.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
extrinsics = np.eye(4, dtype=np.float32)
points = backproject_depth_to_world(depth, intrinsics, extrinsics)
np.testing.assert_allclose(points[0, 0], np.array([0.0, 0.0, 0.75], dtype=np.float32))
def test_estimate_dual_push_target_positions_from_front_rgbd_finds_color_centers() -> None:
rgb = np.zeros((3, 8, 8), dtype=np.float32)
depth = np.full((8, 8), 0.75, dtype=np.float32)
rgb[:, 2:4, 1:3] = np.array([128.0, 128.0, 0.0], dtype=np.float32)[:, None, None] / 255.0
rgb[:, 2:4, 5:7] = np.array([255.0, 128.0, 0.0], dtype=np.float32)[:, None, None] / 255.0
intrinsics = np.array([[-100.0, 0.0, 3.5], [0.0, -100.0, 3.5], [0.0, 0.0, 1.0]], dtype=np.float32)
extrinsics = np.eye(4, dtype=np.float32)
right, left = estimate_dual_push_target_positions_from_front_rgbd(
rgb,
depth,
intrinsics,
extrinsics,
"push the olive and the orange buttons",
)
assert right is not None
assert left is not None
assert right[0] > left[0]
np.testing.assert_allclose(right[2], 0.75, atol=1e-4)
np.testing.assert_allclose(left[2], 0.75, atol=1e-4)
|