feat(pi05): support continuous 6D SE(3) actions (#4132)

* feat(pi05): add continuous 6D SE3 actions

* fix(processors): align rotation 6D with UMI rows

src/lerobot/datasets/compute_stats.py

@@ -18,8 +18,13 @@ from __future__ import annotations
 import logging
 
 import numpy as np
+import torch
 
-from lerobot.processor import RelativeActionsProcessorStep
+from lerobot.processor import (
+    RelativeActionsProcessorStep,
+    relative_action_output_dim,
+    to_relative_actions,
+)
 from lerobot.utils.constants import ACTION, OBS_STATE
 
 from .io_utils import load_image_as_numpy
@@ -660,17 +665,29 @@ def _compute_relative_chunk_batch(
     all_states: np.ndarray,
     chunk_size: int,
     relative_mask: np.ndarray,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: list[list[int]] | None = None,
 ) -> np.ndarray:
     """Vectorised relative-action computation for a batch of start indices.
 
-    Returns an ``(N * chunk_size, action_dim)`` float32 array.
+    Returns an ``(N * chunk_size, model_action_dim)`` float32 array.
     """
     if len(start_indices) == 0:
-        return np.empty((0, all_actions.shape[1]), dtype=np.float32)
+        output_dim = relative_action_output_dim(all_actions.shape[1], pose_representation, se3_pose_groups)
+        return np.empty((0, output_dim), dtype=np.float32)
     offsets = np.arange(chunk_size)
     frame_idx = start_indices[:, None] + offsets[None, :]
     chunks = all_actions[frame_idx].copy()
     states = all_states[start_indices]
+    if pose_representation in {"se3", "se3_6d"}:
+        converted = to_relative_actions(
+            torch.from_numpy(chunks),
+            torch.from_numpy(states),
+            relative_mask.astype(bool).tolist(),
+            pose_representation=pose_representation,
+            se3_pose_groups=se3_pose_groups,
+        )
+        return converted.numpy().reshape(-1, converted.shape[-1])
     mask_dim = len(relative_mask)
     chunks[:, :, :mask_dim] -= states[:, None, :mask_dim] * relative_mask[None, None, :]
     return chunks.reshape(-1, all_actions.shape[1])
@@ -682,6 +699,9 @@ def compute_relative_action_stats(
     chunk_size: int,
     exclude_joints: list[str] | None = None,
     num_workers: int = 0,
+    state_from_action: bool = False,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: list[list[int]] | None = None,
 ) -> dict[str, np.ndarray]:
     """Compute normalization statistics for relative actions over the full dataset.
 
@@ -700,6 +720,9 @@ def compute_relative_action_stats(
         num_workers: Number of parallel threads for computation. Values ≤1
             mean single-threaded. Numpy releases the GIL so threads give
             real parallelism here.
+        state_from_action: Use the current absolute action as state. This is
+            intended for state-less pose datasets where each action row is the
+            synchronized measured robot pose.
 
     Returns:
         Statistics dict with keys "mean", "std", "min", "max", "q01", …, "q99".
@@ -722,7 +745,7 @@ def compute_relative_action_stats(
 
     logging.info("Loading action/state data for relative action stats...")
     all_actions = np.array(hf_dataset[ACTION], dtype=np.float32)
-    all_states = np.array(hf_dataset[OBS_STATE], dtype=np.float32)
+    all_states = all_actions if state_from_action else np.array(hf_dataset[OBS_STATE], dtype=np.float32)
     episode_indices = np.array(hf_dataset["episode_index"])
 
     valid_starts = _get_valid_chunk_starts(episode_indices, chunk_size)
@@ -754,6 +777,8 @@ def compute_relative_action_stats(
                     all_states,
                     chunk_size,
                     relative_mask,
+                    pose_representation,
+                    se3_pose_groups,
                 )
                 for batch in batches
             ]
@@ -762,7 +787,15 @@ def compute_relative_action_stats(
     else:
         for batch in batches:
             running_stats.update(
-                _compute_relative_chunk_batch(batch, all_actions, all_states, chunk_size, relative_mask)
+                _compute_relative_chunk_batch(
+                    batch,
+                    all_actions,
+                    all_states,
+                    chunk_size,
+                    relative_mask,
+                    pose_representation,
+                    se3_pose_groups,
+                )
             )
 
     stats = running_stats.get_statistics()
@@ -777,3 +810,58 @@ def compute_relative_action_stats(
     )
 
     return stats
+
+
+def compute_state_history_stats(
+    hf_dataset,
+    features: dict,
+    history_steps: int,
+    exclude_joints: list[str] | None = None,
+    relative: bool = False,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: list[list[int]] | None = None,
+) -> dict[str, np.ndarray]:
+    """Compute stats for flattened state history synthesized from absolute actions.
+
+    History is left-padded with the first action of each episode, matching dataset
+    boundary padding. When ``relative`` is enabled, every history pose is expressed
+    relative to its newest pose while excluded dimensions remain absolute.
+    """
+    if history_steps < 1:
+        raise ValueError("history_steps must be at least 1")
+    if exclude_joints is None:
+        exclude_joints = []
+
+    actions = np.asarray(hf_dataset[ACTION], dtype=np.float32)
+    episode_indices = np.asarray(hf_dataset["episode_index"])
+    sample_indices = np.arange(len(actions))
+    episode_starts = np.maximum.accumulate(
+        np.where(
+            np.concatenate(([True], episode_indices[1:] != episode_indices[:-1])),
+            sample_indices,
+            0,
+        )
+    )
+    offsets = np.arange(-(history_steps - 1), 1)
+    history_indices = np.maximum(sample_indices[:, None] + offsets[None, :], episode_starts[:, None])
+    history = actions[history_indices].copy()
+
+    if relative:
+        state_dim = actions.shape[-1]
+        names = features.get(ACTION, {}).get("names")
+        mask_step = RelativeActionsProcessorStep(
+            enabled=True,
+            exclude_joints=exclude_joints,
+            action_names=names,
+        )
+        mask = mask_step._build_mask(state_dim)
+        history = to_relative_actions(
+            torch.from_numpy(history),
+            torch.from_numpy(history[:, -1].copy()),
+            mask,
+            pose_representation=pose_representation,
+            se3_pose_groups=se3_pose_groups,
+        ).numpy()
+
+    flattened = history.reshape(len(history), -1)
+    return get_feature_stats(flattened, axis=0, keepdims=False)

src/lerobot/datasets/dataset_tools.py

@@ -54,6 +54,7 @@ from .compute_stats import (
     aggregate_stats,
     compute_episode_stats,
     compute_relative_action_stats,
+    compute_state_history_stats,
 )
 from .dataset_metadata import LeRobotDatasetMetadata
 from .image_writer import write_image
@@ -1566,6 +1567,12 @@ def recompute_stats(
     relative_exclude_joints: list[str] | None = None,
     chunk_size: int = 50,
     num_workers: int = 0,
+    state_from_action: bool = False,
+    state_history_steps: int = 1,
+    relative_state_history: bool = False,
+    relative_state_exclude_joints: list[str] | None = None,
+    relative_pose_representation: str = "componentwise",
+    relative_se3_pose_groups: list[list[int]] | None = None,
 ) -> LeRobotDataset:
     """Recompute stats.json from scratch by iterating all episodes.
 
@@ -1583,6 +1590,16 @@ def recompute_stats(
             ``policy.chunk_size``. Only used when ``relative_action=True``.
         num_workers: Number of parallel threads for relative action stats computation.
             Values ≤1 mean single-threaded. Only used when ``relative_action=True``.
+        state_from_action: Use absolute action rows as synthetic state while
+            computing relative-action stats, and write their absolute statistics
+            under ``observation.state``.
+        state_history_steps: Number of consecutive synthesized state samples.
+        relative_state_history: Express state history relative to its newest pose.
+        relative_state_exclude_joints: State dimensions to retain as absolute.
+        relative_pose_representation: ``componentwise`` for legacy subtraction,
+            ``se3`` for composition with an axis-angle output, or ``se3_6d`` for
+            composition with a continuous two-column rotation output.
+        relative_se3_pose_groups: Six-index xyz+rotation-vector pose groups.
 
     Returns:
         The same dataset with updated stats.
@@ -1606,7 +1623,21 @@ def recompute_stats(
     # (matching what the model sees during training) and skip action in the
     # per-episode pass below.
     relative_action_stats = None
-    if relative_action and ACTION in features and OBS_STATE in features:
+    synthetic_state_stats = None
+    if state_from_action:
+        if ACTION not in features:
+            raise ValueError("state_from_action requires an action feature")
+        synthetic_state_stats = compute_state_history_stats(
+            dataset.hf_dataset,
+            features,
+            history_steps=state_history_steps,
+            exclude_joints=relative_state_exclude_joints,
+            relative=relative_state_history,
+            pose_representation=relative_pose_representation,
+            se3_pose_groups=relative_se3_pose_groups,
+        )
+
+    if relative_action and ACTION in features and (OBS_STATE in features or state_from_action):
         if relative_exclude_joints is None:
             relative_exclude_joints = ["gripper"]
         relative_action_stats = compute_relative_action_stats(
@@ -1615,6 +1646,9 @@ def recompute_stats(
             chunk_size=chunk_size,
             exclude_joints=relative_exclude_joints,
             num_workers=num_workers,
+            state_from_action=state_from_action,
+            pose_representation=relative_pose_representation,
+            se3_pose_groups=relative_se3_pose_groups,
         )
         features_to_compute.pop(ACTION, None)
 
@@ -1654,6 +1688,8 @@ def recompute_stats(
 
     if relative_action_stats is not None:
         new_stats[ACTION] = relative_action_stats
+    if synthetic_state_stats is not None:
+        new_stats[OBS_STATE] = synthetic_state_stats
 
     # Merge: keep existing stats for features we didn't recompute
     if dataset.meta.stats:

src/lerobot/policies/pi05/configuration_pi05.py

@@ -55,6 +55,20 @@ class PI05Config(PreTrainedConfig):
     relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
     # Populated at runtime from dataset metadata by make_policy.
     action_feature_names: list[str] | None = None
+    # ``se3`` uses inv(T_current) @ T_target for each xyz+rotation-vector pose group.
+    # ``se3_6d`` uses the same composition and expands each relative rotation
+    # vector to the continuous first-two-row 6-D rotation representation.
+    # ``componentwise`` preserves the legacy action - state behavior.
+    relative_pose_representation: str = "componentwise"
+    relative_se3_pose_groups: list[list[int]] = field(default_factory=lambda: [list(range(6))])
+
+    # Build proprioception from absolute action samples when the dataset has no
+    # observation.state. With history_steps=2, training samples request t-1 as
+    # well as the normal t..t+chunk_size-1 action targets.
+    state_from_action: bool = False
+    proprioception_history_steps: int = 1
+    use_relative_state_history: bool = False
+    relative_state_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
 
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
@@ -121,6 +135,25 @@ class PI05Config(PreTrainedConfig):
         if self.dtype not in ["bfloat16", "float32"]:
             raise ValueError(f"Invalid dtype: {self.dtype}")
 
+        if self.proprioception_history_steps < 1:
+            raise ValueError("proprioception_history_steps must be at least 1")
+
+        if self.relative_pose_representation not in {"componentwise", "se3", "se3_6d"}:
+            raise ValueError(
+                "relative_pose_representation must be 'componentwise', 'se3', or 'se3_6d', got "
+                f"{self.relative_pose_representation!r}"
+            )
+        for group in self.relative_se3_pose_groups:
+            if len(group) != 6 or len(set(group)) != 6 or any(index < 0 for index in group):
+                raise ValueError(f"Invalid six-index SE(3) pose group: {group}")
+            if self.relative_pose_representation == "se3_6d" and group != list(range(group[0], group[0] + 6)):
+                raise ValueError("se3_6d pose groups must contain six contiguous ascending indices")
+        if self.relative_pose_representation in {"se3", "se3_6d"} and not self.relative_se3_pose_groups:
+            raise ValueError(
+                f"relative_pose_representation={self.relative_pose_representation!r} "
+                "requires relative_se3_pose_groups"
+            )
+
     def validate_features(self) -> None:
         """Validate and set up input/output features."""
         for i in range(self.empty_cameras):
@@ -131,19 +164,54 @@ class PI05Config(PreTrainedConfig):
             )
             self.input_features[key] = empty_camera
 
-        if OBS_STATE not in self.input_features:
-            state_feature = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(self.max_state_dim,),  # Padded to max_state_dim
-            )
-            self.input_features[OBS_STATE] = state_feature
-
         if ACTION not in self.output_features:
             action_feature = PolicyFeature(
                 type=FeatureType.ACTION,
                 shape=(self.max_action_dim,),  # Padded to max_action_dim
             )
             self.output_features[ACTION] = action_feature
+        elif self.relative_pose_representation == "se3_6d":
+            action_feature = self.output_features[ACTION]
+            source_dim = (
+                len(self.action_feature_names)
+                if self.action_feature_names is not None
+                else action_feature.shape[-1]
+            )
+            model_dim = source_dim + 3 * len(self.relative_se3_pose_groups)
+            if action_feature.shape[-1] == source_dim:
+                self.output_features[ACTION] = PolicyFeature(
+                    type=action_feature.type,
+                    shape=(model_dim,),
+                )
+            elif action_feature.shape[-1] != model_dim:
+                raise ValueError(
+                    "se3_6d action feature has incompatible width: "
+                    f"source={source_dim}, expected model width={model_dim}, "
+                    f"got={action_feature.shape[-1]}"
+                )
+            if model_dim > self.max_action_dim:
+                raise ValueError(
+                    f"se3_6d action width {model_dim} exceeds max_action_dim={self.max_action_dim}"
+                )
+
+        if OBS_STATE not in self.input_features:
+            state_shape = (self.max_state_dim,)
+            if self.state_from_action and ACTION in self.output_features:
+                state_shape = self.output_features[ACTION].shape
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=state_shape,
+            )
+            self.input_features[OBS_STATE] = state_feature
+
+        state_dim = self.input_features[OBS_STATE].shape[-1]
+        history_state_dim = state_dim * self.proprioception_history_steps
+        if history_state_dim > self.max_state_dim:
+            raise ValueError(
+                "Flattened proprioception history exceeds max_state_dim: "
+                f"{state_dim} * {self.proprioception_history_steps} = {history_state_dim} > "
+                f"{self.max_state_dim}"
+            )
 
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(
@@ -168,7 +236,8 @@ class PI05Config(PreTrainedConfig):
 
     @property
     def action_delta_indices(self) -> list:
-        return list(range(self.chunk_size))
+        history_prefix = self.proprioception_history_steps - 1 if self.state_from_action else 0
+        return list(range(-history_prefix, self.chunk_size))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/pi05/processor_pi05.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 from copy import deepcopy
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import numpy as np
@@ -36,6 +36,8 @@ from lerobot.processor import (
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
     policy_action_to_transition,
+    relative_action_output_dim,
+    to_relative_actions,
     transition_to_policy_action,
 )
 from lerobot.types import EnvTransition, TransitionKey
@@ -48,6 +50,164 @@ from lerobot.utils.constants import (
 from .configuration_pi05 import PI05Config
 
 
+@ProcessorStepRegistry.register(name="pi05_state_from_action_processor_step")
+@dataclass
+class Pi05StateFromActionProcessorStep(ProcessorStep):
+    """Synthesize proprioception from absolute actions in state-less datasets.
+
+    The dataset loader supplies ``history_steps - 1`` actions before the normal
+    target chunk. Those leading samples and action(t) become state history; only
+    the leading samples are then removed from the action targets.
+    """
+
+    enabled: bool = False
+    history_steps: int = 1
+    _inference_history: torch.Tensor | None = field(default=None, init=False, repr=False)
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        if not self.enabled:
+            return transition
+
+        observation = transition.get(TransitionKey.OBSERVATION, {})
+        observed_state = observation.get(OBS_STATE)
+        if observed_state is not None:
+            # At inference the robot normally provides only the current state and
+            # there is no action target. Build a rolling history in the processor.
+            if transition.get(TransitionKey.ACTION) is None and observed_state.ndim == 2:
+                if self._inference_history is None:
+                    self._inference_history = observed_state.unsqueeze(1).repeat(1, self.history_steps, 1)
+                else:
+                    self._inference_history = torch.cat(
+                        [self._inference_history[:, 1:], observed_state.unsqueeze(1)], dim=1
+                    )
+                new_transition = transition.copy()
+                new_observation = dict(observation)
+                new_observation[OBS_STATE] = self._inference_history.clone()
+                new_transition[TransitionKey.OBSERVATION] = new_observation
+                return new_transition
+            return transition
+
+        action = transition.get(TransitionKey.ACTION)
+        if action is None:
+            raise ValueError("Cannot synthesize PI0.5 state without action")
+        if action.ndim != 3:
+            raise ValueError(f"Expected batched action chunks with shape (B, T, D), got {action.shape}")
+        if action.shape[1] < self.history_steps:
+            raise ValueError(
+                f"Action chunk has {action.shape[1]} steps, fewer than history_steps={self.history_steps}"
+            )
+
+        new_transition = transition.copy()
+        new_observation = dict(observation)
+        state = action[:, : self.history_steps].clone()
+        if self.history_steps == 1:
+            state = state[:, 0]
+        new_observation[OBS_STATE] = state
+        new_transition[TransitionKey.OBSERVATION] = new_observation
+        new_transition[TransitionKey.ACTION] = action[:, self.history_steps - 1 :]
+        return new_transition
+
+    def get_config(self) -> dict[str, Any]:
+        return {"enabled": self.enabled, "history_steps": self.history_steps}
+
+    def reset(self) -> None:
+        self._inference_history = None
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@ProcessorStepRegistry.register(name="pi05_flatten_state_history_processor_step")
+@dataclass
+class Pi05FlattenStateHistoryProcessorStep(ProcessorStep):
+    """Optionally relativize raw state history, then flatten it for PI0.5."""
+
+    history_steps: int = 1
+    max_state_dim: int = 32
+    relative: bool = False
+    exclude_joints: list[str] = field(default_factory=list)
+    state_names: list[str] | None = None
+    pose_representation: str = "componentwise"
+    se3_pose_groups: list[list[int]] = field(default_factory=list)
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        observation = transition.get(TransitionKey.OBSERVATION, {})
+        state = observation.get(OBS_STATE)
+        if state is None:
+            raise ValueError("State is required for PI05")
+        if self.history_steps == 1 and state.ndim == 2:
+            state = state.unsqueeze(1)
+        if state.ndim != 3 or state.shape[1] != self.history_steps:
+            raise ValueError(
+                f"Expected state history with shape (B, {self.history_steps}, D), got {state.shape}"
+            )
+
+        processed_state = state.clone()
+        if self.relative:
+            mask_step = RelativeActionsProcessorStep(
+                enabled=True,
+                exclude_joints=self.exclude_joints,
+                action_names=self.state_names,
+            )
+            processed_state = to_relative_actions(
+                state,
+                state[:, -1],
+                mask_step._build_mask(state.shape[-1]),
+                pose_representation=self.pose_representation,
+                se3_pose_groups=self.se3_pose_groups,
+            )
+
+        flattened_dim = processed_state.shape[1] * processed_state.shape[2]
+        if flattened_dim > self.max_state_dim:
+            raise ValueError(
+                f"Flattened state history has {flattened_dim} dimensions, above max_state_dim={self.max_state_dim}"
+            )
+
+        new_transition = transition.copy()
+        new_observation = dict(observation)
+        new_observation[OBS_STATE] = processed_state.flatten(start_dim=1)
+        new_transition[TransitionKey.OBSERVATION] = new_observation
+        return new_transition
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "history_steps": self.history_steps,
+            "max_state_dim": self.max_state_dim,
+            "relative": self.relative,
+            "exclude_joints": self.exclude_joints,
+            "state_names": self.state_names,
+            "pose_representation": self.pose_representation,
+            "se3_pose_groups": self.se3_pose_groups,
+        }
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        transformed = deepcopy(features)
+        for feature_group in transformed.values():
+            state_feature = feature_group.get(OBS_STATE)
+            if state_feature is not None:
+                state_dim = state_feature.shape[-1]
+                if self.relative:
+                    source_dim = len(self.state_names) if self.state_names is not None else state_dim
+                    model_dim = relative_action_output_dim(
+                        source_dim,
+                        self.pose_representation,
+                        self.se3_pose_groups,
+                    )
+                    if state_dim == source_dim:
+                        state_dim = model_dim
+                    elif state_dim != model_dim:
+                        raise ValueError(
+                            f"Expected source/model state width {source_dim}/{model_dim}, got {state_dim}"
+                        )
+                state_dim *= self.history_steps
+                feature_group[OBS_STATE] = PolicyFeature(type=state_feature.type, shape=(state_dim,))
+        return transformed
+
+
 @ProcessorStepRegistry.register(name="pi05_prepare_state_tokenizer_processor_step")
 @dataclass
 class Pi05PrepareStateTokenizerProcessorStep(ProcessorStep):
@@ -133,13 +293,28 @@ def make_pi05_pre_post_processors(
         enabled=config.use_relative_actions,
         exclude_joints=getattr(config, "relative_exclude_joints", []),
         action_names=getattr(config, "action_feature_names", None),
+        pose_representation=config.relative_pose_representation,
+        se3_pose_groups=config.relative_se3_pose_groups,
     )
 
     # OpenPI order: raw → relative → normalize → model → unnormalize → absolute
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
         AddBatchDimensionProcessorStep(),
+        Pi05StateFromActionProcessorStep(
+            enabled=config.state_from_action,
+            history_steps=config.proprioception_history_steps,
+        ),
         relative_step,
+        Pi05FlattenStateHistoryProcessorStep(
+            history_steps=config.proprioception_history_steps,
+            max_state_dim=config.max_state_dim,
+            relative=config.use_relative_state_history,
+            exclude_joints=config.relative_state_exclude_joints,
+            state_names=config.action_feature_names,
+            pose_representation=config.relative_pose_representation,
+            se3_pose_groups=config.relative_se3_pose_groups,
+        ),
         # NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
         # because the tokenizer step expects normalized state in [-1, 1] range for discretization
         NormalizerProcessorStep(

src/lerobot/processor/__init__.py

@@ -89,8 +89,15 @@ from .policy_robot_bridge import (
 from .relative_action_processor import (
     AbsoluteActionsProcessorStep,
     RelativeActionsProcessorStep,
+    relative_action_output_dim,
+    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,
 )
 from .rename_processor import RenameObservationsProcessorStep, rename_stats
 from .tokenizer_processor import ActionTokenizerProcessorStep, TokenizerProcessorStep
@@ -135,6 +142,15 @@ __all__ = [
     "make_default_robot_observation_processor",
     "AbsoluteActionsProcessorStep",
     "RelativeActionsProcessorStep",
+    "relative_action_output_dim",
+    "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",
     "MapDeltaActionToRobotActionStep",
     "MapTensorToDeltaActionDictStep",
     "NewLineTaskProcessorStep",
@@ -168,8 +184,6 @@ __all__ = [
     "transition_to_batch",
     "TransitionKey",
     "TruncatedProcessorStep",
-    "to_absolute_actions",
-    "to_relative_actions",
     "UnnormalizerProcessorStep",
     "VanillaObservationProcessorStep",
 ]

src/lerobot/processor/relative_action_processor.py

@@ -21,7 +21,7 @@ from torch import Tensor
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import OBS_STATE
+from lerobot.utils.constants import ACTION, OBS_STATE
 
 from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
 from .pipeline import ProcessorStep, ProcessorStepRegistry
@@ -34,57 +34,399 @@ __all__ = [
     "AbsoluteActionsProcessorStep",
     "to_relative_actions",
     "to_absolute_actions",
+    "to_relative_se3_pose",
+    "to_absolute_se3_pose",
+    "to_relative_se3_pose_6d",
+    "to_absolute_se3_pose_6d",
+    "rotation_6d_to_rotvec",
+    "rotvec_to_rotation_6d",
+    "relative_action_output_dim",
 ]
 
 
-def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
+def _rotvec_to_quaternion(rotvec: Tensor) -> Tensor:
-    """Convert absolute actions to relative: relative = action - state (for masked dims).
+    angle = torch.linalg.vector_norm(rotvec, dim=-1, keepdim=True)
+    angle_sq = angle.square()
+    small_scale = 0.5 - angle_sq / 48.0 + angle_sq.square() / 3840.0
+    scale = torch.where(angle > 1e-6, torch.sin(angle / 2.0) / angle.clamp_min(1e-12), small_scale)
+    return torch.cat((torch.cos(angle / 2.0), rotvec * scale), dim=-1)
+
+
+def _quaternion_to_rotvec(quaternion: Tensor) -> Tensor:
+    quaternion = quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+    quaternion = quaternion * torch.where(quaternion[..., :1] < 0, -1.0, 1.0)
+    vector = quaternion[..., 1:]
+    sin_half_angle = torch.linalg.vector_norm(vector, dim=-1, keepdim=True)
+    angle = 2.0 * torch.atan2(sin_half_angle, quaternion[..., :1].clamp_min(0.0))
+    small_scale = 2.0 + sin_half_angle.square() / 3.0
+    scale = torch.where(
+        sin_half_angle > 1e-6,
+        angle / sin_half_angle.clamp_min(1e-12),
+        small_scale,
+    )
+    return vector * scale
+
+
+def _quaternion_multiply(left: Tensor, right: Tensor) -> Tensor:
+    left_w, left_xyz = left[..., :1], left[..., 1:]
+    right_w, right_xyz = right[..., :1], right[..., 1:]
+    return torch.cat(
+        (
+            left_w * right_w - (left_xyz * right_xyz).sum(dim=-1, keepdim=True),
+            left_w * right_xyz + right_w * left_xyz + torch.linalg.cross(left_xyz, right_xyz, dim=-1),
+        ),
+        dim=-1,
+    )
+
+
+def _quaternion_conjugate(quaternion: Tensor) -> Tensor:
+    return torch.cat((quaternion[..., :1], -quaternion[..., 1:]), dim=-1)
+
+
+def _quaternion_rotate(quaternion: Tensor, vector: Tensor) -> Tensor:
+    quaternion_xyz = quaternion[..., 1:]
+    uv = torch.linalg.cross(quaternion_xyz, vector, dim=-1)
+    uuv = torch.linalg.cross(quaternion_xyz, uv, dim=-1)
+    return vector + 2.0 * (quaternion[..., :1] * uv + uuv)
+
+
+def _quaternion_to_matrix(quaternion: Tensor) -> Tensor:
+    quaternion = quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+    w, x, y, z = quaternion.unbind(-1)
+    two_s = 2.0
+    return torch.stack(
+        (
+            1.0 - two_s * (y * y + z * z),
+            two_s * (x * y - z * w),
+            two_s * (x * z + y * w),
+            two_s * (x * y + z * w),
+            1.0 - two_s * (x * x + z * z),
+            two_s * (y * z - x * w),
+            two_s * (x * z - y * w),
+            two_s * (y * z + x * w),
+            1.0 - two_s * (x * x + y * y),
+        ),
+        dim=-1,
+    ).reshape(quaternion.shape[:-1] + (3, 3))
+
+
+def _matrix_to_quaternion(matrix: Tensor) -> Tensor:
+    """Convert proper rotation matrices to normalized ``[w, x, y, z]`` quaternions."""
+    if matrix.shape[-2:] != (3, 3):
+        raise ValueError(f"Rotation matrices must have shape (..., 3, 3), got {matrix.shape}")
+
+    m00 = matrix[..., 0, 0]
+    m01 = matrix[..., 0, 1]
+    m02 = matrix[..., 0, 2]
+    m10 = matrix[..., 1, 0]
+    m11 = matrix[..., 1, 1]
+    m12 = matrix[..., 1, 2]
+    m20 = matrix[..., 2, 0]
+    m21 = matrix[..., 2, 1]
+    m22 = matrix[..., 2, 2]
+
+    # Each row is a quaternion candidate scaled by the magnitude of its
+    # best-conditioned component. Selecting the largest component avoids the
+    # trace singularity at rotations close to pi.
+    q_abs = torch.sqrt(
+        torch.clamp(
+            torch.stack(
+                (
+                    1.0 + m00 + m11 + m22,
+                    1.0 + m00 - m11 - m22,
+                    1.0 - m00 + m11 - m22,
+                    1.0 - m00 - m11 + m22,
+                ),
+                dim=-1,
+            ),
+            min=0.0,
+        )
+    )
+    quat_by_rijk = torch.stack(
+        (
+            torch.stack((q_abs[..., 0].square(), m21 - m12, m02 - m20, m10 - m01), dim=-1),
+            torch.stack((m21 - m12, q_abs[..., 1].square(), m10 + m01, m02 + m20), dim=-1),
+            torch.stack((m02 - m20, m10 + m01, q_abs[..., 2].square(), m12 + m21), dim=-1),
+            torch.stack((m10 - m01, m02 + m20, m12 + m21, q_abs[..., 3].square()), dim=-1),
+        ),
+        dim=-2,
+    )
+    candidates = quat_by_rijk / (2.0 * q_abs[..., :, None].clamp_min(0.1))
+    best = torch.nn.functional.one_hot(q_abs.argmax(dim=-1), num_classes=4).to(dtype=matrix.dtype)
+    quaternion = (candidates * best[..., :, None]).sum(dim=-2)
+    return quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+
+
+def rotvec_to_rotation_6d(rotvec: Tensor) -> Tensor:
+    """Encode an axis-angle rotation as the first two rotation-matrix rows."""
+    matrix = _quaternion_to_matrix(_rotvec_to_quaternion(rotvec))
+    return matrix[..., :2, :].reshape(matrix.shape[:-2] + (6,))
+
+
+def rotation_6d_to_rotvec(rotation_6d: Tensor) -> Tensor:
+    """Decode two predicted 3-D vectors into an axis-angle rotation.
+
+    Gram-Schmidt orthonormalization follows the continuous 6-D rotation
+    representation. Degenerate predictions fail closed instead of producing an
+    invalid physical rotation.
+    """
+    if rotation_6d.shape[-1] != 6:
+        raise ValueError(f"6-D rotations must have six values, got {rotation_6d.shape}")
+    first = rotation_6d[..., :3]
+    second = rotation_6d[..., 3:]
+    first_norm = torch.linalg.vector_norm(first, dim=-1, keepdim=True)
+    first_unit = first / first_norm.clamp_min(1e-12)
+    second_orthogonal = second - (first_unit * second).sum(dim=-1, keepdim=True) * first_unit
+    second_norm = torch.linalg.vector_norm(second_orthogonal, dim=-1, keepdim=True)
+    if bool(torch.any(first_norm <= 1e-8)) or bool(torch.any(second_norm <= 1e-8)):
+        raise ValueError("Cannot decode a degenerate 6-D rotation prediction")
+    second_unit = second_orthogonal / second_norm
+    third_unit = torch.linalg.cross(first_unit, second_unit, dim=-1)
+    matrix = torch.stack((first_unit, second_unit, third_unit), dim=-2)
+    return _quaternion_to_rotvec(_matrix_to_quaternion(matrix))
+
+
+def to_relative_se3_pose(target_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Encode a pose as ``inv(T_reference) @ T_target``.
+
+    Poses use ``[x, y, z, rx, ry, rz]`` with an axis-angle rotation vector.
+    The relative translation is therefore expressed in the reference EE frame.
+    """
+    if target_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
+        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
+    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
+    target_quaternion = _rotvec_to_quaternion(target_pose[..., 3:])
+    inverse_reference_quaternion = _quaternion_conjugate(reference_quaternion)
+    relative_translation = _quaternion_rotate(
+        inverse_reference_quaternion, target_pose[..., :3] - reference_pose[..., :3]
+    )
+    relative_quaternion = _quaternion_multiply(inverse_reference_quaternion, target_quaternion)
+    return torch.cat((relative_translation, _quaternion_to_rotvec(relative_quaternion)), dim=-1)
+
+
+def to_absolute_se3_pose(relative_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Decode a pose with ``T_target = T_reference @ T_relative``."""
+    if relative_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
+        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
+    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
+    relative_quaternion = _rotvec_to_quaternion(relative_pose[..., 3:])
+    target_translation = reference_pose[..., :3] + _quaternion_rotate(
+        reference_quaternion, relative_pose[..., :3]
+    )
+    target_quaternion = _quaternion_multiply(reference_quaternion, relative_quaternion)
+    return torch.cat((target_translation, _quaternion_to_rotvec(target_quaternion)), dim=-1)
+
+
+def to_relative_se3_pose_6d(target_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Encode ``inv(T_reference) @ T_target`` as xyz plus continuous 6-D rotation."""
+    relative_pose = to_relative_se3_pose(target_pose, reference_pose)
+    return torch.cat((relative_pose[..., :3], rotvec_to_rotation_6d(relative_pose[..., 3:])), dim=-1)
+
+
+def to_absolute_se3_pose_6d(relative_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Decode xyz plus continuous 6-D rotation with ``T_target = T_reference @ T_relative``."""
+    if relative_pose.shape[-1] != 9:
+        raise ValueError("6-D encoded SE(3) poses must have nine values: xyz plus rotation-6D")
+    relative_rotvec_pose = torch.cat(
+        (relative_pose[..., :3], rotation_6d_to_rotvec(relative_pose[..., 3:])), dim=-1
+    )
+    return to_absolute_se3_pose(relative_rotvec_pose, reference_pose)
+
+
+def _broadcast_reference(actions: Tensor, state: Tensor) -> Tensor:
+    if state.device != actions.device or state.dtype != actions.dtype:
+        state = state.to(device=actions.device, dtype=actions.dtype)
+    if actions.ndim == state.ndim + 1:
+        state = state.unsqueeze(-2)
+    return state
+
+
+def _validate_se3_pose_groups(
+    pose_representation: str,
+    se3_pose_groups: Sequence[Sequence[int]] | None,
+    mask: Sequence[bool],
+    action_dim: int,
+) -> list[list[int]]:
+    if pose_representation not in {"componentwise", "se3", "se3_6d"}:
+        raise ValueError(
+            f"Unsupported pose_representation={pose_representation!r}; expected "
+            "'componentwise', 'se3', or 'se3_6d'"
+        )
+    if pose_representation == "componentwise":
+        return []
+    if not se3_pose_groups:
+        raise ValueError(
+            f"pose_representation={pose_representation!r} requires at least one six-index se3_pose_group"
+        )
+
+    normalized_groups: list[list[int]] = []
+    used_indices: set[int] = set()
+    for raw_group in se3_pose_groups:
+        group = [int(index) for index in raw_group]
+        if len(group) != 6:
+            raise ValueError(f"Each SE(3) pose group must contain six indices, got {group}")
+        if len(set(group)) != 6 or any(index < 0 or index >= action_dim for index in group):
+            raise ValueError(f"Invalid SE(3) pose group for action_dim={action_dim}: {group}")
+        if pose_representation == "se3_6d" and group != list(range(group[0], group[0] + 6)):
+            raise ValueError("se3_6d pose groups must contain six contiguous ascending indices")
+        if any(index >= len(mask) for index in group):
+            raise ValueError(f"SE(3) pose group lies outside the relative mask: {group}")
+        if used_indices.intersection(group):
+            raise ValueError(f"SE(3) pose groups must not overlap: {group}")
+        group_mask = [bool(mask[index]) for index in group]
+        if any(group_mask) and not all(group_mask):
+            raise ValueError(f"An SE(3) pose group must be wholly relative or wholly absolute: {group}")
+        used_indices.update(group)
+        if all(group_mask):
+            normalized_groups.append(group)
+    return normalized_groups
+
+
+def relative_action_output_dim(
+    source_dim: int,
+    pose_representation: str,
+    se3_pose_groups: Sequence[Sequence[int]] | None,
+) -> int:
+    """Return the model-space action width for a source action width."""
+    if pose_representation != "se3_6d":
+        return source_dim
+    groups = se3_pose_groups or []
+    return source_dim + 3 * len(groups)
+
+
+def _expand_se3_6d_actions(
+    actions: Tensor,
+    state: Tensor,
+    groups: Sequence[Sequence[int]],
+) -> Tensor:
+    group_by_start = {group[0]: list(group) for group in groups}
+    grouped_indices = {index for group in groups for index in group}
+    parts: list[Tensor] = []
+    for index in range(actions.shape[-1]):
+        group = group_by_start.get(index)
+        if group is not None:
+            parts.append(to_relative_se3_pose_6d(actions[..., group], state[..., group]))
+        elif index not in grouped_indices:
+            parts.append(actions[..., index : index + 1])
+    return torch.cat(parts, dim=-1)
+
+
+def _collapse_se3_6d_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    groups: Sequence[Sequence[int]],
+) -> Tensor:
+    source_dim = len(mask)
+    expected_dim = relative_action_output_dim(source_dim, "se3_6d", groups)
+    if actions.shape[-1] != expected_dim:
+        raise ValueError(
+            f"Expected se3_6d action width {expected_dim} for source width {source_dim}, "
+            f"got {actions.shape[-1]}"
+        )
+    group_by_start = {group[0]: list(group) for group in groups}
+    grouped_indices = {index for group in groups for index in group}
+    parts: list[Tensor] = []
+    cursor = 0
+    for index in range(source_dim):
+        group = group_by_start.get(index)
+        if group is not None:
+            parts.append(to_absolute_se3_pose_6d(actions[..., cursor : cursor + 9], state[..., group]))
+            cursor += 9
+        elif index not in grouped_indices:
+            value = actions[..., cursor : cursor + 1]
+            if mask[index]:
+                value = value + state[..., index : index + 1]
+            parts.append(value)
+            cursor += 1
+    if cursor != actions.shape[-1]:
+        raise RuntimeError(f"Consumed {cursor} action values from width {actions.shape[-1]}")
+    return torch.cat(parts, dim=-1)
+
+
+def to_relative_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    *,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: Sequence[Sequence[int]] | None = None,
+) -> Tensor:
+    """Convert absolute actions to a configured relative representation.
+
+    Component-wise mode computes ``action - state``. SE(3) modes compute
+    ``inv(T_state) @ T_action`` for each configured pose group. ``se3_6d``
+    replaces each three-value relative rotation vector with its continuous
+    six-value encoding, increasing the output width by three per pose group.
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
+    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, actions.shape[-1])
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    if state.device != actions.device or state.dtype != actions.dtype:
+    state = _broadcast_reference(actions, state)
-        state = state.to(device=actions.device, dtype=actions.dtype)
+    component_mask = mask_t.clone()
-    state_offset = state[..., :dims] * mask_t
+    for group in groups:
-    if actions.ndim == 3:
+        component_mask[group] = 0
-        state_offset = state_offset.unsqueeze(-2)
+    state_offset = state[..., :dims] * component_mask
     actions = actions.clone()
     actions[..., :dims] -= state_offset
+    if pose_representation == "se3_6d":
+        return _expand_se3_6d_actions(actions, state, groups)
+    for group in groups:
+        actions[..., group] = to_relative_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
-def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
+def to_absolute_actions(
-    """Convert relative actions back to absolute: absolute = relative + state (for masked dims).
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    *,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: Sequence[Sequence[int]] | None = None,
+) -> Tensor:
+    """Convert relative actions back to absolute actions.
+
+    Component-wise mode computes ``relative + state``. SE(3) mode computes
+    ``T_state @ T_relative`` for each configured pose group.
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
+    source_dim = len(mask)
+    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, source_dim)
+    state = _broadcast_reference(actions, state)
+    if pose_representation == "se3_6d":
+        return _collapse_se3_6d_actions(actions, state, mask, groups)
+
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    if state.device != actions.device or state.dtype != actions.dtype:
+    state = _broadcast_reference(actions, state)
-        state = state.to(device=actions.device, dtype=actions.dtype)
+    component_mask = mask_t.clone()
-    state_offset = state[..., :dims] * mask_t
+    for group in groups:
-    if actions.ndim == 3:
+        component_mask[group] = 0
-        state_offset = state_offset.unsqueeze(-2)
+    state_offset = state[..., :dims] * component_mask
     actions = actions.clone()
     actions[..., :dims] += state_offset
+    for group in groups:
+        actions[..., group] = to_absolute_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
 @ProcessorStepRegistry.register("relative_actions_processor")
 @dataclass
 class RelativeActionsProcessorStep(ProcessorStep):
-    """Converts absolute actions to relative actions (action -= state) for masked dimensions.
+    """Converts absolute actions to the configured relative representation.
 
     Mirrors OpenPI's DeltaActions transform. Applied during preprocessing so the model
     trains on relative offsets instead of absolute positions.
@@ -101,7 +443,10 @@ class RelativeActionsProcessorStep(ProcessorStep):
     enabled: bool = False
     exclude_joints: list[str] = field(default_factory=list)
     action_names: list[str] | None = None
+    pose_representation: str = "componentwise"
+    se3_pose_groups: list[list[int]] = field(default_factory=list)
     _last_state: torch.Tensor | None = field(default=None, init=False, repr=False)
+    _last_mask: list[bool] | None = field(default=None, init=False, repr=False)
 
     def _build_mask(self, action_dim: int) -> list[bool]:
         if not self.exclude_joints or self.action_names is None:
@@ -126,37 +471,78 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
+        # State history has shape (B, H, D). Relative actions are referenced to
+        # the newest proprioceptive state, not the whole history tensor.
+        reference_state = state[:, -1] if state is not None and state.ndim == 3 else state
+
         # Always cache state for the paired AbsoluteActionsProcessorStep
-        if state is not None:
+        if reference_state is not None:
-            self._last_state = state
+            self._last_state = reference_state
+            self._last_mask = self._build_mask(reference_state.shape[-1])
 
         if not self.enabled:
             return transition
 
         new_transition = transition.copy()
         action = new_transition.get(TransitionKey.ACTION)
-        if action is None or state is None:
+        if action is None or reference_state is None:
             return new_transition
 
-        mask = self._build_mask(action.shape[-1])
+        mask = self._last_mask or self._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_relative_actions(action, state, mask)
+        new_transition[TransitionKey.ACTION] = to_relative_actions(
+            action,
+            reference_state,
+            mask,
+            pose_representation=self.pose_representation,
+            se3_pose_groups=self.se3_pose_groups,
+        )
         return new_transition
 
     def get_cached_state(self) -> torch.Tensor | None:
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
+    def get_cached_mask(self) -> list[bool] | None:
+        """Return the source-space mask cached with the latest state."""
+        return self._last_mask
+
+    def reset(self) -> None:
+        """Drop the inference reference so it cannot leak between sessions."""
+        self._last_state = None
+        self._last_mask = None
+
     def get_config(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,
             "exclude_joints": self.exclude_joints,
             "action_names": self.action_names,
+            "pose_representation": self.pose_representation,
+            "se3_pose_groups": self.se3_pose_groups,
         }
 
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
+        if not self.enabled or self.pose_representation != "se3_6d":
+            return features
+        transformed = {feature_type: dict(feature_group) for feature_type, feature_group in features.items()}
+        for feature_group in transformed.values():
+            action_feature = feature_group.get(ACTION)
+            if action_feature is None:
+                continue
+            source_dim = len(self.action_names) if self.action_names is not None else action_feature.shape[-1]
+            model_dim = relative_action_output_dim(source_dim, self.pose_representation, self.se3_pose_groups)
+            if action_feature.shape[-1] == source_dim:
+                feature_group[ACTION] = PolicyFeature(
+                    type=action_feature.type,
+                    shape=(model_dim,),
+                )
+            elif action_feature.shape[-1] != model_dim:
+                raise ValueError(
+                    f"Expected source/model action width {source_dim}/{model_dim}, "
+                    f"got {action_feature.shape[-1]}"
+                )
+        return transformed
 
 
 @ProcessorStepRegistry.register("absolute_actions_processor")
@@ -198,8 +584,16 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
         if action is None:
             return new_transition
 
-        mask = self.relative_step._build_mask(action.shape[-1])
+        mask = self.relative_step.get_cached_mask()
-        new_transition[TransitionKey.ACTION] = to_absolute_actions(action, cached_state, mask)
+        if mask is None:
+            mask = self.relative_step._build_mask(cached_state.shape[-1])
+        new_transition[TransitionKey.ACTION] = to_absolute_actions(
+            action,
+            cached_state,
+            mask,
+            pose_representation=self.relative_step.pose_representation,
+            se3_pose_groups=self.relative_step.se3_pose_groups,
+        )
         return new_transition
 
     def get_config(self) -> dict[str, Any]:

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -325,6 +325,8 @@ class RecomputeStatsConfig(OperationConfig):
     relative_exclude_joints: list[str] | None = None
     chunk_size: int = 50
     num_workers: int = 0
+    relative_pose_representation: str = "componentwise"
+    relative_se3_pose_groups: list[list[int]] | None = None
     overwrite: bool = False
 
 
@@ -698,6 +700,8 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
         relative_exclude_joints=cfg.operation.relative_exclude_joints,
         chunk_size=cfg.operation.chunk_size,
         num_workers=cfg.operation.num_workers,
+        relative_pose_representation=cfg.operation.relative_pose_representation,
+        relative_se3_pose_groups=cfg.operation.relative_se3_pose_groups,
     )
 
     logging.info(f"Stats written to {dataset.root}")

src/lerobot/scripts/lerobot_train.py

@@ -343,6 +343,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                 "enabled": True,
                 "exclude_joints": getattr(active_cfg, "relative_exclude_joints", []),
                 "action_names": getattr(active_cfg, "action_feature_names", None),
+                "pose_representation": getattr(active_cfg, "relative_pose_representation", "componentwise"),
+                "se3_pose_groups": getattr(active_cfg, "relative_se3_pose_groups", []),
             }
             postprocessor_overrides["absolute_actions_processor"] = {"enabled": True}
         processor_kwargs["preprocessor_overrides"] = preprocessor_overrides

tests/processor/test_pi05_state_from_action.py

@@ -0,0 +1,321 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from math import pi
+
+import numpy as np
+import pytest
+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,
+)
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.processor_pi05 import (  # noqa: E402
+    Pi05FlattenStateHistoryProcessorStep,
+    Pi05StateFromActionProcessorStep,
+)
+from lerobot.processor.relative_action_processor import (  # noqa: E402
+    AbsoluteActionsProcessorStep,
+    RelativeActionsProcessorStep,
+)
+from lerobot.types import TransitionKey  # noqa: E402
+from lerobot.utils.constants import ACTION, OBS_STATE  # noqa: E402
+
+
+def _transition(action: torch.Tensor | None, state: torch.Tensor | None = None) -> dict:
+    observation = {} if state is None else {OBS_STATE: state}
+    return {
+        TransitionKey.OBSERVATION: observation,
+        TransitionKey.ACTION: action,
+        TransitionKey.REWARD: None,
+        TransitionKey.DONE: None,
+        TransitionKey.TRUNCATED: None,
+        TransitionKey.COMPLEMENTARY_DATA: {},
+    }
+
+
+def test_pi05_config_requests_action_history_prefix():
+    config = PI05Config(
+        device="cpu",
+        chunk_size=4,
+        n_action_steps=4,
+        state_from_action=True,
+        proprioception_history_steps=2,
+    )
+
+    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)
+
+    result = step(_transition(action))
+
+    torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], action[:, :2])
+    torch.testing.assert_close(result[TransitionKey.ACTION], action[:, 1:])
+
+
+def test_relative_actions_use_newest_state_in_history_and_roundtrip():
+    state_history = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
+    absolute = torch.tensor([[[3.0, 30.0], [4.0, 40.0]]])
+    relative_step = RelativeActionsProcessorStep(enabled=True)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+
+    relative = relative_step(_transition(absolute, state_history))
+    expected = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
+    torch.testing.assert_close(relative[TransitionKey.ACTION], expected)
+
+    recovered = absolute_step(_transition(relative[TransitionKey.ACTION]))
+    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)
+
+    result = step(_transition(torch.zeros(1, 2, 2), state_history))
+
+    torch.testing.assert_close(
+        result[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+    )
+
+
+def test_state_history_can_be_relative_with_absolute_gripper():
+    state_history = torch.tensor([[[1.0, 10.0, 0.2], [3.0, 20.0, 0.4]]])
+    step = Pi05FlattenStateHistoryProcessorStep(
+        history_steps=2,
+        max_state_dim=6,
+        relative=True,
+        exclude_joints=["gripper"],
+        state_names=["x", "y", "gripper_width"],
+    )
+
+    result = step(_transition(torch.zeros(1, 2, 3), state_history))
+
+    torch.testing.assert_close(
+        result[TransitionKey.OBSERVATION][OBS_STATE],
+        torch.tensor([[-2.0, -10.0, 0.2, 0.0, 0.0, 0.4]]),
+    )
+
+
+def test_state_history_can_use_se3_composition_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=14,
+        relative=True,
+        exclude_joints=["gripper"],
+        state_names=["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
+        pose_representation="se3",
+        se3_pose_groups=[list(range(6))],
+    )
+
+    result = step(_transition(torch.zeros(1, 2, 7), state_history))
+
+    expected = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.2, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.4]])
+    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)
+
+    first = step(_transition(None, torch.tensor([[1.0, 2.0]])))
+    second = step(_transition(None, torch.tensor([[3.0, 4.0]])))
+    torch.testing.assert_close(
+        first[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [1.0, 2.0]]])
+    )
+    torch.testing.assert_close(
+        second[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
+    )
+
+    step.reset()
+    reset = step(_transition(None, torch.tensor([[5.0, 6.0]])))
+    torch.testing.assert_close(
+        reset[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[5.0, 6.0], [5.0, 6.0]]])
+    )
+
+
+def test_flatten_state_history_checks_max_state_dim():
+    step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=3)
+
+    with pytest.raises(ValueError, match="above max_state_dim"):
+        step(_transition(torch.zeros(1, 2, 2), torch.zeros(1, 2, 2)))
+
+
+def test_relative_stats_can_use_absolute_action_as_state():
+    actions = np.asarray([[0.0, 0.0], [1.0, 2.0], [2.0, 4.0], [3.0, 6.0]], dtype=np.float32)
+    dataset = {"action": actions, "episode_index": np.zeros(4, dtype=np.int64)}
+    features = {"action": {"shape": [2], "names": ["x", "y"]}}
+
+    stats = compute_relative_action_stats(
+        dataset,
+        features,
+        chunk_size=2,
+        state_from_action=True,
+    )
+
+    np.testing.assert_allclose(stats["mean"], [0.5, 1.0])
+
+
+def test_relative_state_history_stats_match_processor_representation():
+    actions = np.asarray(
+        [[0.0, 0.1], [1.0, 0.2], [3.0, 0.3]],
+        dtype=np.float32,
+    )
+    dataset = {"action": actions, "episode_index": np.zeros(3, dtype=np.int64)}
+    features = {"action": {"shape": [2], "names": ["x", "gripper_width"]}}
+
+    stats = compute_state_history_stats(
+        dataset,
+        features,
+        history_steps=2,
+        exclude_joints=["gripper"],
+        relative=True,
+    )
+
+    expected = np.asarray([[0.0, 0.1, 0.0, 0.1], [-1.0, 0.1, 0.0, 0.2], [-2.0, 0.2, 0.0, 0.3]])
+    np.testing.assert_allclose(stats["mean"], expected.mean(axis=0))
+
+
+def test_se3_relative_action_stats_use_reference_frame():
+    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"],
+        }
+    }
+
+    stats = compute_relative_action_stats(
+        dataset,
+        features,
+        chunk_size=2,
+        exclude_joints=["gripper"],
+        state_from_action=True,
+        pose_representation="se3",
+        se3_pose_groups=[list(range(6))],
+    )
+
+    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

@@ -0,0 +1,144 @@
+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_rotation_6d_uses_umi_first_two_rows():
+    source = torch.tensor([[0.0, 0.0, math.pi / 2]], dtype=torch.float64)
+
+    encoded = rotvec_to_rotation_6d(source)
+
+    expected = torch.tensor([[0.0, -1.0, 0.0, 1.0, 0.0, 0.0]], dtype=torch.float64)
+    torch.testing.assert_close(encoded, expected, atol=1e-7, rtol=1e-7)
+    torch.testing.assert_close(rotation_6d_to_rotvec(expected), source, atol=1e-7, rtol=1e-7)
+
+
+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))