import abc
from typing import TYPE_CHECKING, Dict, Literal, Optional, Set, Tuple, Union, overload
import numpy as np
import pinocchio as pin
from typing_extensions import Self
from robokit import CONFIG
from robokit.types import ArrayLike
try:
import torch # torch is optional
_TORCH_AVAILABLE = True
except ImportError:
_TORCH_AVAILABLE = False
try:
import warp as wp # warp is optional
_WARP_AVAILABLE = True
except ImportError:
_WARP_AVAILABLE = False
if TYPE_CHECKING:
from robokit.lie.pinocchio_se3 import PinocchioSE3
from robokit.lie.torch_se3 import TorchSE3
from robokit.lie.warp_se3 import WarpSE3
# Configuration for SE3 implementations
# Maps array_type to (default_backend, supported_backends)
SE3_BACKEND_CONFIG: Dict[str, Tuple[str, Set[str]]] = {
"numpy": ("pinocchio", {"pinocchio"}),
"torch": (CONFIG.default_torch_compute_backend, {"torch", "warp"}),
"warp": ("warp", {"warp"}),
}
[docs]
class SE3(abc.ABC):
"""SE(3) representation.
Internal parameterization: [x, y, z, qw, qx, qy, qz].
Tangent parameterization: [vx, vy, vz, omega_x, omega_y, omega_z].
"""
@staticmethod
def _infer_array_type(
param: Union[pin.SE3, ArrayLike],
array_type: Optional[Literal["numpy", "torch", "warp"]] = None,
) -> Literal["numpy", "torch", "warp"]:
if array_type is None:
if isinstance(param, (np.ndarray, pin.SE3)):
array_type = "numpy"
elif _TORCH_AVAILABLE and isinstance(param, torch.Tensor):
array_type = "torch"
elif _WARP_AVAILABLE and isinstance(param, wp.array):
array_type = "warp"
else:
raise ValueError(
f"Cannot infer array_type from type: {type(param)}. Expected pin.SE3, numpy.ndarray, torch.Tensor, or wp.array."
)
return array_type
@staticmethod
def _get_se3_class(
array_type: Literal["numpy", "torch", "warp"],
) -> Union["type[PinocchioSE3]", "type[TorchSE3]", "type[WarpSE3]"]:
if array_type == "numpy":
from robokit.lie.pinocchio_se3 import PinocchioSE3
return PinocchioSE3
elif array_type == "torch":
from robokit.lie.torch_se3 import TorchSE3
return TorchSE3
elif array_type == "warp":
from robokit.lie.warp_se3 import WarpSE3
return WarpSE3
else:
raise ValueError(f"Unsupported array_type: {array_type}")
@staticmethod
def _validate_backends(
array_type: Literal["numpy", "torch", "warp"],
compute_backend: Optional[Literal["pinocchio", "torch", "warp"]],
) -> Literal["pinocchio", "torch", "warp"]:
default_backend, supported_backends = SE3_BACKEND_CONFIG[array_type]
if compute_backend is None:
return default_backend # type: ignore
if compute_backend not in supported_backends:
raise ValueError(
f"array_type='{array_type}' cannot use compute_backend='{compute_backend}'. Valid options: {supported_backends}"
)
return compute_backend
@overload
def __new__(
cls, se3_like: "wp.array", array_type: Literal["warp"], compute_backend: Optional[Literal["warp"]] = ...
) -> "WarpSE3": ...
@overload
def __new__(
cls,
se3_like: "torch.Tensor",
array_type: Optional[Literal["torch"]] = ...,
compute_backend: Optional[Literal["torch", "warp"]] = ...,
) -> "TorchSE3": ...
@overload
def __new__(
cls,
se3_like: Union[pin.SE3, np.ndarray],
array_type: Optional[Literal["numpy"]] = ...,
compute_backend: Optional[Literal["pinocchio"]] = ...,
) -> "PinocchioSE3": ...
def __new__(
cls,
se3_like: Union[pin.SE3, ArrayLike],
array_type: Optional[Literal["numpy", "torch", "warp"]] = None,
compute_backend: Optional[Literal["pinocchio", "torch", "warp"]] = None,
) -> "SE3":
# If se3_like is already a SE3 instance, return it directly
if isinstance(se3_like, SE3):
return se3_like
# If called directly on SE3, determine the array_type and return appropriate subclass
if cls is SE3:
array_type = cls._infer_array_type(se3_like, array_type)
compute_backend = cls._validate_backends(array_type, compute_backend)
se3_cls = cls._get_se3_class(array_type)
return se3_cls.__new__(se3_cls, se3_like, array_type, compute_backend) # type: ignore
else:
# Called on subclass, use normal instantiation
return super().__new__(cls)
@property
@abc.abstractmethod
def xyz_wxyz(self) -> Union[np.ndarray, "torch.Tensor"]:
"""Internal parameterization of SE(3) as [x, y, z, qw, qx, qy, qz]"""
...
@property
@abc.abstractmethod
def xyz(self) -> Union[np.ndarray, "torch.Tensor"]:
"""Translation part of SE(3) as [x, y, z]"""
...
@property
@abc.abstractmethod
def quat_wxyz(self) -> Union[np.ndarray, "torch.Tensor"]:
"""Quaternion part of SE(3) as [qw, qx, qy, qz]"""
...
@abc.abstractmethod
def as_matrix(self) -> Union[np.ndarray, "torch.Tensor"]: ...
@staticmethod
def exp(
log_transform: Union[np.ndarray, "torch.Tensor"],
array_type: Optional[Literal["numpy", "torch", "warp"]] = None,
compute_backend: Optional[Literal["pinocchio", "torch", "warp"]] = None,
) -> "SE3":
array_type = SE3._infer_array_type(log_transform, array_type)
compute_backend = SE3._validate_backends(array_type, compute_backend)
se3_cls = SE3._get_se3_class(array_type)
return se3_cls.exp(log_transform, array_type, compute_backend) # type: ignore
@abc.abstractmethod
def inverse(self) -> Self: ...
[docs]
@abc.abstractmethod
def log(self) -> Union[np.ndarray, "torch.Tensor"]:
"""The log representation in [v, omega] format, linear first, then angular."""
...
[docs]
@abc.abstractmethod
def adjoint(self) -> Union[np.ndarray, "torch.Tensor"]:
"""The 6x6 adjoint matrix Ad_T = [[R, skew(t) @ R], [0, R]] mapping twists as v' = Ad_T @ v with [v, omega] ordering (linear first, then angular)."""
...
# NOTE jlog for TorchSE3/WarpSE3 is not implemented yet
# @abc.abstractmethod
def jlog(self) -> Union[np.ndarray, "torch.Tensor"]: ...
def __getitem__(self, key) -> Self:
raise NotImplementedError
@abc.abstractmethod
def __repr__(self) -> str: ...
def __str__(self) -> str:
return self.__repr__()
@abc.abstractmethod
def __mul__(self, other: "SE3") -> Self: ...
@abc.abstractmethod
def __rmul__(self, other: "SE3") -> Self: ...
def __matmul__(self, other: "SE3") -> Self:
return self.__mul__(other)
def __rmatmul__(self, other: "SE3") -> Self:
return self.__rmul__(other)