feat(pi05): add continuous 6D SE3 actions

tests/processor/test_pi05_state_from_action.py

@@ -22,6 +22,7 @@ import torch
 
 pytest.importorskip("transformers")
 
+from lerobot.configs import FeatureType, PolicyFeature  # noqa: E402
 from lerobot.datasets.compute_stats import (  # noqa: E402
     compute_relative_action_stats,
     compute_state_history_stats,
@@ -36,7 +37,7 @@ from lerobot.processor.relative_action_processor import (  # noqa: E402
     RelativeActionsProcessorStep,
 )
 from lerobot.types import TransitionKey  # noqa: E402
-from lerobot.utils.constants import OBS_STATE  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
 
 
 def _transition(action: torch.Tensor | None, state: torch.Tensor | None = None) -> dict:
@@ -63,6 +64,27 @@ def test_pi05_config_requests_action_history_prefix():
     assert config.action_delta_indices == [-1, 0, 1, 2, 3]
 
 
+def test_pi05_config_accepts_se3_6d_action_and_state_with_two_step_history():
+    names = ["x", "y", "z", "rx", "ry", "rz", "gripper_width"]
+    config = PI05Config(
+        device="cpu",
+        use_relative_actions=True,
+        state_from_action=True,
+        proprioception_history_steps=2,
+        use_relative_state_history=True,
+        relative_pose_representation="se3_6d",
+        action_feature_names=names,
+        output_features={
+            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+        },
+    )
+
+    config.validate_features()
+
+    assert config.output_features[ACTION].shape == (10,)
+    assert config.input_features[OBS_STATE].shape == (10,)
+
+
 def test_state_from_action_extracts_history_and_preserves_target_horizon():
     action = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3)
     step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
@@ -87,6 +109,16 @@ def test_relative_actions_use_newest_state_in_history_and_roundtrip():
     torch.testing.assert_close(recovered[TransitionKey.ACTION], absolute)
 
 
+def test_relative_action_reference_is_reset_between_inference_sessions():
+    step = RelativeActionsProcessorStep(enabled=True)
+    step(_transition(None, torch.tensor([[1.0, 2.0]])))
+
+    step.reset()
+
+    assert step.get_cached_state() is None
+    assert step.get_cached_mask() is None
+
+
 def test_flatten_state_history_preserves_chronological_order():
     state_history = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
     step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=4)
@@ -136,6 +168,27 @@ def test_state_history_can_use_se3_composition_with_absolute_gripper():
     torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], expected, atol=1e-6, rtol=1e-6)
 
 
+def test_state_history_can_use_se3_6d_rotation_with_absolute_gripper():
+    state_history = torch.tensor(
+        [[[0.0, 1.0, 0.0, 0.0, 0.0, pi / 2, 0.2], [0.0, 0.0, 0.0, 0.0, 0.0, pi / 2, 0.4]]]
+    )
+    step = Pi05FlattenStateHistoryProcessorStep(
+        history_steps=2,
+        max_state_dim=20,
+        relative=True,
+        exclude_joints=["gripper"],
+        state_names=["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
+        pose_representation="se3_6d",
+        se3_pose_groups=[list(range(6))],
+    )
+
+    result = step(_transition(torch.zeros(1, 2, 7), state_history))
+
+    identity_6d = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0]
+    expected = torch.tensor([[1.0, 0.0, 0.0, *identity_6d, 0.2, 0.0, 0.0, 0.0, *identity_6d, 0.4]])
+    torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], expected, atol=1e-6, rtol=1e-6)
+
+
 def test_inference_state_history_is_rolled_and_reset():
     step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
 
@@ -225,3 +278,44 @@ def test_se3_relative_action_stats_use_reference_frame():
 
     np.testing.assert_allclose(stats["mean"][:3], [0.5, 0.0, 0.0], atol=1e-6)
     np.testing.assert_allclose(stats["mean"][6], 0.25, atol=1e-6)
+
+
+def test_se3_6d_stats_expand_action_and_state_history():
+    actions = np.asarray(
+        [
+            [0.0, 0.0, 0.0, 0.0, 0.0, pi / 2, 0.2],
+            [0.0, 1.0, 0.0, 0.0, 0.0, pi / 2, 0.3],
+        ],
+        dtype=np.float32,
+    )
+    dataset = {"action": actions, "episode_index": np.zeros(2, dtype=np.int64)}
+    features = {
+        "action": {
+            "shape": [7],
+            "names": ["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
+        }
+    }
+
+    action_stats = compute_relative_action_stats(
+        dataset,
+        features,
+        chunk_size=2,
+        exclude_joints=["gripper"],
+        state_from_action=True,
+        pose_representation="se3_6d",
+        se3_pose_groups=[list(range(6))],
+    )
+    state_stats = compute_state_history_stats(
+        dataset,
+        features,
+        history_steps=2,
+        exclude_joints=["gripper"],
+        relative=True,
+        pose_representation="se3_6d",
+        se3_pose_groups=[list(range(6))],
+    )
+
+    assert action_stats["mean"].shape == (10,)
+    assert state_stats["mean"].shape == (20,)
+    np.testing.assert_allclose(action_stats["mean"][:3], [0.5, 0.0, 0.0], atol=1e-6)
+    np.testing.assert_allclose(action_stats["mean"][9], 0.25, atol=1e-6)

tests/processor/test_se3_relative_actions.py

@@ -4,10 +4,14 @@ 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))]
@@ -68,3 +72,63 @@ def test_se3_pose_group_cannot_be_partially_relative():
             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))