diff --git a/src/lerobot/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py index 14901cc4c..ef8de0b49 100644 --- a/src/lerobot/datasets/compute_stats.py +++ b/src/lerobot/datasets/compute_stats.py @@ -20,7 +20,11 @@ import logging import numpy as np import torch -from lerobot.processor import RelativeActionsProcessorStep, to_relative_actions +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 @@ -666,26 +670,24 @@ def _compute_relative_chunk_batch( ) -> 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 == "se3": - return ( - 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, - ) - .numpy() - .reshape(-1, all_actions.shape[1]) + 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]) diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 1cc0b23fe..78f883cc9 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -1596,8 +1596,9 @@ def recompute_stats( 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 or - ``se3`` for ``inv(T_current) @ T_target`` pose composition. + 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: diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index a14e002a0..1c419baa6 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/src/lerobot/policies/pi05/configuration_pi05.py @@ -56,6 +56,8 @@ class PI05Config(PreTrainedConfig): # 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 two-column 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))]) @@ -136,16 +138,21 @@ class PI05Config(PreTrainedConfig): 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"}: + if self.relative_pose_representation not in {"componentwise", "se3", "se3_6d"}: raise ValueError( - "relative_pose_representation must be either 'componentwise' or 'se3', got " + "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" and not self.relative_se3_pose_groups: - raise ValueError("relative_pose_representation='se3' requires relative_se3_pose_groups") + 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.""" @@ -163,6 +170,29 @@ class PI05Config(PreTrainedConfig): 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,) diff --git a/src/lerobot/policies/pi05/processor_pi05.py b/src/lerobot/policies/pi05/processor_pi05.py index 094e7e5ac..f1b69fddf 100644 --- a/src/lerobot/policies/pi05/processor_pi05.py +++ b/src/lerobot/policies/pi05/processor_pi05.py @@ -36,6 +36,7 @@ from lerobot.processor import ( TokenizerProcessorStep, UnnormalizerProcessorStep, policy_action_to_transition, + relative_action_output_dim, to_relative_actions, transition_to_policy_action, ) @@ -143,12 +144,6 @@ class Pi05FlattenStateHistoryProcessorStep(ProcessorStep): f"Expected state history with shape (B, {self.history_steps}, D), got {state.shape}" ) - flattened_dim = state.shape[1] * 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}" - ) - processed_state = state.clone() if self.relative: mask_step = RelativeActionsProcessorStep( @@ -164,6 +159,12 @@ class Pi05FlattenStateHistoryProcessorStep(ProcessorStep): 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) @@ -187,8 +188,22 @@ class Pi05FlattenStateHistoryProcessorStep(ProcessorStep): transformed = deepcopy(features) for feature_group in transformed.values(): state_feature = feature_group.get(OBS_STATE) - if state_feature is not None and self.history_steps > 1: - state_dim = state_feature.shape[-1] * self.history_steps + 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 diff --git a/src/lerobot/processor/__init__.py b/src/lerobot/processor/__init__.py index 32e4b892d..24f52c821 100644 --- a/src/lerobot/processor/__init__.py +++ b/src/lerobot/processor/__init__.py @@ -89,10 +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 @@ -137,10 +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", diff --git a/src/lerobot/processor/relative_action_processor.py b/src/lerobot/processor/relative_action_processor.py index 06b949b06..09d3cf6a8 100644 --- a/src/lerobot/processor/relative_action_processor.py +++ b/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 @@ -36,6 +36,11 @@ __all__ = [ "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", ] @@ -85,6 +90,102 @@ def _quaternion_rotate(quaternion: Tensor, vector: Tensor) -> Tensor: 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 two rotation-matrix columns.""" + matrix = _quaternion_to_matrix(_rotvec_to_quaternion(rotvec)) + return torch.cat((matrix[..., :, 0], matrix[..., :, 1]), dim=-1) + + +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=-1) + 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``. @@ -116,6 +217,22 @@ def to_absolute_se3_pose(relative_pose: Tensor, reference_pose: Tensor) -> Tenso 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) @@ -130,14 +247,17 @@ def _validate_se3_pose_groups( mask: Sequence[bool], action_dim: int, ) -> list[list[int]]: - if pose_representation not in {"componentwise", "se3"}: + if pose_representation not in {"componentwise", "se3", "se3_6d"}: raise ValueError( - f"Unsupported pose_representation={pose_representation!r}; expected 'componentwise' or 'se3'" + 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("pose_representation='se3' requires at least one six-index se3_pose_group") + 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() @@ -147,6 +267,8 @@ def _validate_se3_pose_groups( 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): @@ -160,6 +282,68 @@ def _validate_se3_pose_groups( 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, @@ -170,8 +354,10 @@ def to_relative_actions( ) -> Tensor: """Convert absolute actions to a configured relative representation. - Component-wise mode computes ``action - state``. SE(3) mode computes - ``inv(T_state) @ T_action`` for each configured pose group. + 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). @@ -190,6 +376,8 @@ def to_relative_actions( 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 @@ -213,7 +401,12 @@ def to_absolute_actions( 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]) + 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 @@ -253,6 +446,7 @@ class RelativeActionsProcessorStep(ProcessorStep): 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: @@ -284,6 +478,7 @@ class RelativeActionsProcessorStep(ProcessorStep): # Always cache state for the paired AbsoluteActionsProcessorStep if reference_state is not None: self._last_state = reference_state + self._last_mask = self._build_mask(reference_state.shape[-1]) if not self.enabled: return transition @@ -293,7 +488,7 @@ class RelativeActionsProcessorStep(ProcessorStep): 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, reference_state, @@ -307,6 +502,15 @@ class RelativeActionsProcessorStep(ProcessorStep): """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, @@ -319,7 +523,26 @@ class RelativeActionsProcessorStep(ProcessorStep): 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") @@ -361,7 +584,9 @@ 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() + if mask is None: + mask = self.relative_step._build_mask(cached_state.shape[-1]) new_transition[TransitionKey.ACTION] = to_absolute_actions( action, cached_state, diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 6e8458523..ad4212924 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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 diff --git a/tests/processor/test_pi05_state_from_action.py b/tests/processor/test_pi05_state_from_action.py index 053e57969..03ec2ed9d 100644 --- a/tests/processor/test_pi05_state_from_action.py +++ b/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) diff --git a/tests/processor/test_se3_relative_actions.py b/tests/processor/test_se3_relative_actions.py index 75639bef9..64dde8145 100644 --- a/tests/processor/test_se3_relative_actions.py +++ b/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))