import abc
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Set, Tuple, Union, overload
import numpy as np
import yourdfpy
from robokit import CONFIG
from robokit.lie.se3 import SE3
from robokit.robo.robot_spec import RobotSpec
from robokit.types import ArrayLike
if TYPE_CHECKING:
from robokit.robo.numpy_robot import NumpyRobot
from robokit.robo.torch_robot import TorchRobot
from robokit.robo.warp_robot import WarpRobot
logger = logging.getLogger("robokit")
# Configuration for Robot implementations
# Maps array_type to supported_backends
ROBOT_SUPPORTED_BACKENDS: Dict[str, Set[str]] = {"numpy": {"pinocchio"}, "torch": {"torch", "warp"}, "warp": {"warp"}}
def _get_default_backend(array_type: Literal["numpy", "torch", "warp"]) -> str:
if array_type == "numpy":
return "pinocchio"
elif array_type == "torch":
return CONFIG.default_torch_compute_backend
elif array_type == "warp":
return "warp"
else:
raise ValueError(f"Unsupported array_type: {array_type}")
[docs]
@dataclass
class RobotState(abc.ABC):
spec: RobotSpec
@abc.abstractmethod
def set_configuration(self, q: ArrayLike, T_world_base: Optional[SE3] = None): ...
@abc.abstractmethod
def get_T_world_link(self, link_index: int) -> SE3: ...
[docs]
@abc.abstractmethod
def get_link_jacobian(
self, link_index: int, reference_frame: Literal["body", "spatial"] = "body"
) -> Union[ArrayLike, Tuple[ArrayLike, ArrayLike]]:
"""
For fixed-base robots, returns only the joint Jacobian.
For floating-base robots (when T_world_base is provided), returns both joint and base Jacobians.
"""
...
def __repr__(self) -> str:
return f"RobotState(q={self.q}, T_world_base={self.T_world_base})"
def __str__(self) -> str:
return self.__repr__()
[docs]
class Robot(abc.ABC):
"""
Robot model with multiple backends.
Examples:
>>> from robot_descriptions.loaders.yourdfpy import load_robot_description
>>> robot = Robot.load(load_robot_description("panda_description"), array_type="numpy")
>>> q = np.array([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0, 0.0], dtype=np.float32)
>>> state = robot.state(q=q)
>>> state = robot.forward_kinematics(state)
>>> ee_pose = state.get_T_world_link(robot.link_names.index("panda_hand"))
>>> expected_pose = np.array([0.0608599, 0.0, 0.7637312, 0.0382045, 0.9192637, 0.3807715, 0.0922340])
>>> np.allclose(ee_pose.xyz_wxyz, expected_pose, atol=1e-4)
True
>>> import torch
>>> robot = Robot.load(load_robot_description("panda_description"), array_type="torch")
>>> q = torch.tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0, 0.0], dtype=torch.float32)
>>> state = robot.state(q=q)
>>> state = robot.forward_kinematics(state)
>>> ee_pose = state.get_T_world_link(robot.link_names.index("panda_hand"))
>>> torch.allclose(ee_pose.xyz_wxyz, torch.tensor([0.0608599, 0.0, 0.7637312, 0.0382045, 0.9192637, 0.3807715, 0.0922340], dtype=torch.float32), atol=1e-4)
True
"""
spec: RobotSpec
@staticmethod
def _get_robot_class(
array_type: Literal["numpy", "torch", "warp"],
) -> Union["type[NumpyRobot]", "type[TorchRobot]", "type[WarpRobot]"]:
# fmt: off
if array_type == "numpy":
from robokit.robo.numpy_robot import NumpyRobot
return NumpyRobot
elif array_type == "torch":
from robokit.robo.torch_robot import TorchRobot
return TorchRobot
elif array_type == "warp":
from robokit.robo.warp_robot import WarpRobot
return WarpRobot
else:
raise ValueError(f"Unsupported array_type: {array_type}")
# fmt: on
@staticmethod
def _validate_backends(
array_type: Literal["numpy", "torch", "warp"],
compute_backend: Optional[Literal["pinocchio", "torch", "warp"]],
) -> Literal["pinocchio", "torch", "warp"]:
if compute_backend is None:
return _get_default_backend(array_type) # type: ignore
if compute_backend not in ROBOT_SUPPORTED_BACKENDS[array_type]:
raise ValueError(
f"array_type='{array_type}' cannot use compute_backend='{compute_backend}'. Valid options: {ROBOT_SUPPORTED_BACKENDS[array_type]}"
)
return compute_backend
@overload
def __new__(
cls, spec: RobotSpec, array_type: Literal["numpy"], compute_backend: Optional[Literal["pinocchio"]] = ...
) -> "NumpyRobot": ...
@overload
def __new__(
cls, spec: RobotSpec, array_type: Literal["torch"], compute_backend: Optional[Literal["torch", "warp"]] = ...
) -> "TorchRobot": ...
@overload
def __new__(
cls, spec: RobotSpec, array_type: Literal["warp"], compute_backend: Optional[Literal["warp"]] = ...
) -> "WarpRobot": ...
def __new__(
cls,
spec: RobotSpec,
array_type: Literal["numpy", "torch", "warp"] = "numpy",
compute_backend: Optional[Literal["pinocchio", "torch", "warp"]] = None,
) -> "Robot":
if cls is Robot:
compute_backend = cls._validate_backends(array_type, compute_backend)
robot_cls = cls._get_robot_class(array_type)
return robot_cls.__new__(robot_cls, spec, array_type, compute_backend) # type: ignore
else:
# Called on subclass, use normal instantiation
return super().__new__(cls)
# fmt: off
@overload
@staticmethod
def load(robot_description_or_path: Union[str, Path, yourdfpy.URDF], array_type: Literal["numpy"] = ..., compute_backend: Optional[Literal["pinocchio"]] = ..., load_meshes: bool = ..., mesh_dir: Optional[Union[str, Path]] = ...) -> "NumpyRobot": ...
@overload
@staticmethod
def load(robot_description_or_path: Union[str, Path, yourdfpy.URDF], array_type: Literal["torch"] = ..., compute_backend: Optional[Literal["torch", "warp"]] = ..., load_meshes: bool = ..., mesh_dir: Optional[Union[str, Path]] = ...) -> "TorchRobot": ...
@overload
@staticmethod
def load(robot_description_or_path: Union[str, Path, yourdfpy.URDF], array_type: Literal["warp"] = ..., compute_backend: Optional[Literal["warp"]] = ..., load_meshes: bool = ..., mesh_dir: Optional[Union[str, Path]] = ...) -> "WarpRobot": ...
# fmt: on
@staticmethod
def load(
robot_description_or_path: Union[str, Path, yourdfpy.URDF],
array_type: Literal["numpy", "torch", "warp"] = "numpy",
compute_backend: Optional[Literal["pinocchio", "torch", "warp"]] = None,
load_meshes: bool = False,
mesh_dir: Optional[Union[str, Path]] = None,
) -> "Robot":
return Robot(RobotSpec.parse(robot_description_or_path, load_meshes, mesh_dir), array_type, compute_backend) # type: ignore
@abc.abstractmethod
def state(self) -> RobotState: ...
@abc.abstractmethod
def forward_kinematics(self, state: RobotState) -> RobotState: ...
@abc.abstractmethod
def compute_motion_subspace(self, state: RobotState) -> RobotState: ...
def __repr__(self) -> str:
return (
f"Robot(name={self.spec.name}, num_actuated_joints={self.spec.num_actuated_joints}, num_links={self.spec.num_links})\n"
+ self.spec.kinematic_tree_str
)
def __str__(self) -> str:
return self.__repr__()
@property
def num_joints(self) -> int:
return self.spec.num_joints
@property
def num_actuated_joints(self) -> int:
return self.spec.num_actuated_joints
@property
def num_nonfixed_joints(self) -> int:
return self.spec.num_nonfixed_joints
@property
def num_dofs(self) -> int:
return self.spec.num_dofs
@property
def num_links(self) -> int:
return self.spec.num_links
@property
def link_names(self) -> List[str]:
return self.spec.link_names
@property
def joint_names(self) -> List[str]:
return self.spec.joint_names
@property
def actuated_joint_names(self) -> List[str]:
return self.spec.actuated_joint_names
@property
def actuated_joint_limits(self) -> np.ndarray:
return self.spec.actuated_joint_limits
@property
def midrange_q(self) -> np.ndarray:
return self.spec.midrange_q
@property
def zero_q(self) -> np.ndarray:
return self.spec.zero_q
__all__ = ["Robot", "RobotState"]