mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 10:16:09 +00:00
Pepijn
135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
import math
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from lerobot.processor.relative_action_processor import (
|
|
rotation_6d_to_rotvec,
|
|
rotvec_to_rotation_6d,
|
|
to_absolute_actions,
|
|
to_absolute_se3_pose,
|
|
to_absolute_se3_pose_6d,
|
|
to_relative_actions,
|
|
to_relative_se3_pose,
|
|
to_relative_se3_pose_6d,
|
|
)
|
|
|
|
POSE_GROUP = [list(range(6))]
|
|
|
|
|
|
def test_se3_translation_is_expressed_in_reference_frame():
|
|
reference = torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, math.pi / 2]])
|
|
target = torch.tensor([[1.0, 3.0, 3.0, 0.0, 0.0, math.pi / 2]])
|
|
|
|
relative = to_relative_se3_pose(target, reference)
|
|
|
|
torch.testing.assert_close(relative, torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), atol=1e-6, rtol=1e-6)
|
|
|
|
|
|
def test_se3_pose_roundtrip_for_batched_chunks():
|
|
torch.manual_seed(0)
|
|
reference = torch.randn(4, 6)
|
|
reference[:, 3:] *= 0.8
|
|
target = torch.randn(4, 11, 6)
|
|
target[..., 3:] *= 0.8
|
|
|
|
relative = to_relative_se3_pose(target, reference.unsqueeze(1))
|
|
recovered = to_absolute_se3_pose(relative, reference.unsqueeze(1))
|
|
|
|
torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
|
|
|
|
|
|
def test_mixed_se3_pose_and_absolute_gripper_roundtrip():
|
|
reference = torch.tensor([[0.2, -0.1, 0.4, 0.1, 0.2, -0.3, 0.06]])
|
|
target = torch.tensor([[[0.3, 0.2, 0.5, -0.2, 0.1, 0.4, 0.03], [0.1, -0.3, 0.2, 0.5, -0.1, 0.2, 0.05]]])
|
|
mask = [True, True, True, True, True, True, False]
|
|
|
|
relative = to_relative_actions(
|
|
target,
|
|
reference,
|
|
mask,
|
|
pose_representation="se3",
|
|
se3_pose_groups=POSE_GROUP,
|
|
)
|
|
recovered = to_absolute_actions(
|
|
relative,
|
|
reference,
|
|
mask,
|
|
pose_representation="se3",
|
|
se3_pose_groups=POSE_GROUP,
|
|
)
|
|
|
|
torch.testing.assert_close(relative[..., 6], target[..., 6])
|
|
torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
|
|
|
|
|
|
def test_se3_pose_group_cannot_be_partially_relative():
|
|
with pytest.raises(ValueError, match="wholly relative or wholly absolute"):
|
|
to_relative_actions(
|
|
torch.zeros(1, 7),
|
|
torch.zeros(1, 7),
|
|
[True, True, True, False, False, False, False],
|
|
pose_representation="se3",
|
|
se3_pose_groups=POSE_GROUP,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"rotvec",
|
|
[
|
|
[0.0, 0.0, 0.0],
|
|
[0.2, -0.5, 0.8],
|
|
[math.pi - 1e-4, 0.0, 0.0],
|
|
],
|
|
)
|
|
def test_rotation_6d_roundtrip(rotvec):
|
|
source = torch.tensor([rotvec], dtype=torch.float64)
|
|
|
|
recovered = rotation_6d_to_rotvec(rotvec_to_rotation_6d(source))
|
|
|
|
torch.testing.assert_close(recovered, source, atol=2e-6, rtol=2e-6)
|
|
|
|
|
|
def test_se3_6d_pose_roundtrip_for_batched_chunks():
|
|
torch.manual_seed(1)
|
|
reference = torch.randn(4, 6)
|
|
reference[:, 3:] *= 0.8
|
|
target = torch.randn(4, 11, 6)
|
|
target[..., 3:] *= 0.8
|
|
|
|
relative = to_relative_se3_pose_6d(target, reference.unsqueeze(1))
|
|
recovered = to_absolute_se3_pose_6d(relative, reference.unsqueeze(1))
|
|
|
|
assert relative.shape == (4, 11, 9)
|
|
torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
|
|
|
|
|
|
def test_mixed_se3_6d_pose_and_absolute_gripper_roundtrip():
|
|
reference = torch.tensor([[0.2, -0.1, 0.4, 0.1, 0.2, -0.3, 0.06]])
|
|
target = torch.tensor([[[0.3, 0.2, 0.5, -0.2, 0.1, 0.4, 0.03], [0.1, -0.3, 0.2, 0.5, -0.1, 0.2, 0.05]]])
|
|
mask = [True, True, True, True, True, True, False]
|
|
|
|
relative = to_relative_actions(
|
|
target,
|
|
reference,
|
|
mask,
|
|
pose_representation="se3_6d",
|
|
se3_pose_groups=POSE_GROUP,
|
|
)
|
|
recovered = to_absolute_actions(
|
|
relative,
|
|
reference,
|
|
mask,
|
|
pose_representation="se3_6d",
|
|
se3_pose_groups=POSE_GROUP,
|
|
)
|
|
|
|
assert relative.shape == (1, 2, 10)
|
|
torch.testing.assert_close(relative[..., 9], target[..., 6])
|
|
torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
|
|
|
|
|
|
def test_rotation_6d_rejects_degenerate_prediction():
|
|
with pytest.raises(ValueError, match="degenerate"):
|
|
rotation_6d_to_rotvec(torch.zeros(1, 6))
|