se3labs-interface 0.0.1a1__tar.gz

se3labs_interface-0.0.1a1/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: se3labs-interface
+Version: 0.0.1a1
+Summary: SE3 Labs SDK shared interface
+Requires-Python: >=3.10
+Requires-Dist: numpy>=2.2
+Requires-Dist: jaxtyping>=0.3
+Requires-Dist: beartype>=0.22
+Requires-Dist: grpcio>=1.83.1
+Requires-Dist: protobuf>=7.35.1
+Requires-Dist: pillow>=10
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pyright; extra == "dev"

se3labs_interface-0.0.1a1/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "se3labs-interface"
+version = "0.0.1a1"
+description = "SE3 Labs SDK shared interface"
+requires-python = ">=3.10"
+dependencies = [
+    "numpy>=2.2",
+    "jaxtyping>=0.3",
+    "beartype>=0.22",
+    "grpcio>=1.83.1",
+    "protobuf>=7.35.1",
+    "pillow>=10",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "pyright",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+"se3labs.interface.wire" = ["*.proto", "*.pyi"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.pyright]
+include = ["src", "tests"]
+exclude = ["**/__pycache__", "**/*_pb2.py", "**/*_pb2_grpc.py"]
+typeCheckingMode = "standard"

se3labs_interface-0.0.1a1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/__init__.py

@@ -0,0 +1,23 @@
+from . import arm, camera, job, policy, station, task
+from .arm import (
+    MotorizedArm_ArmType,
+    MotorizedArm_GripperType,
+    MotorizedArmSpec,
+    MotorizedArm_Observation,
+    MotorizedArm_JointPose,
+    MotorizedArm_ActionChunk,
+)
+from .camera import Camera_DeviceType, Camera_DeviceSpec, Camera_Observation
+from .station import (
+    Station_BaseSpec,
+    Station_Spec,
+    Bimanual_Station_Spec,
+    Single_Arm_Station_Spec,
+    Station_Observation,
+    Station_JointPose,
+    Station_ActionChunk,
+)
+from .task       import Task_ScoreKind, Task_ScoreChoice, Task_ScoreItem, Task_ScoreAnswers, Task_ScoreSheet, Task_Spec
+from .policy     import Policy
+from .job        import JobRequest
+from .repository import Bimanual_YAM_Station

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/arm.py

@@ -0,0 +1,169 @@
+"""Motorized arm hardware spec and the per-arm observation / action containers.
+
+Conventions (all arrays are ``np.float32``):
+
+* joint angles in radians, joint velocities in rad/s, joint torques in N*m,
+  joint temperatures in degrees Celsius.
+* ``gripper_pos`` is a normalized open fraction in ``[0, 1]`` where ``0`` is
+  fully closed and ``1`` is fully open. The station maps this to the physical
+  gripper travel, so policies never see device units.
+* Action targets are **absolute** joint positions.
+"""
+import numpy as np
+import jaxtyping as jt
+from beartype import beartype
+from enum import Enum
+from dataclasses import dataclass
+
+from ..util.validation import check_hint
+
+
+class MotorizedArm_ArmType(Enum):
+    YAM_6DoF = "yam_6dof"
+    """A YAM 6-DoF arm.
+
+    Reference: https://i2rt.com/products/yam-6-dof-arm
+    """
+
+    @property
+    def n_dof(self) -> int:
+        return _ARM_N_DOF[self]
+
+
+_ARM_N_DOF: dict[MotorizedArm_ArmType, int] = {
+    MotorizedArm_ArmType.YAM_6DoF: 6,
+}
+
+
+class MotorizedArm_GripperType(Enum):
+    I2RT_linear_gripper_1DoF = "i2rt_linear_gripper_1dof"
+    """I2RT linear gripper with 1 DoF.
+
+    Reference: https://i2rt.com/products/yam-gripper-1
+    """
+
+    @property
+    def n_dof(self) -> int:
+        return _GRIPPER_N_DOF[self]
+
+
+_GRIPPER_N_DOF: dict[MotorizedArm_GripperType, int] = {
+    MotorizedArm_GripperType.I2RT_linear_gripper_1DoF: 1,
+}
+
+
+@dataclass(frozen=True)
+class MotorizedArmSpec:
+    arm         : MotorizedArm_ArmType
+    gripper     : MotorizedArm_GripperType
+
+    refresh_rate: int
+    """Rate at which observations are published, unit: [Hz]."""
+
+    control_rate: int
+    """Rate at which joint position targets are sent to the motors, unit: [Hz].
+
+    Action chunks emitted at a lower rate are interpolated up to this rate by the station.
+    """
+
+    @property
+    def n_arm_dof(self) -> int:
+        return self.arm.n_dof
+
+    @property
+    def n_gripper_dof(self) -> int:
+        return self.gripper.n_dof
+
+
+@jt.jaxtyped(typechecker=beartype)
+@dataclass(frozen=True)
+class MotorizedArm_Observation:
+    """One observation tick of a single arm. Every channel is always populated.
+
+    Construction enforces dtype / rank and that every joint channel shares ``n_arm_dof``;
+    ``validate`` binds ``n_arm_dof`` / ``n_gripper_dof`` to a concrete ``MotorizedArmSpec``.
+    """
+
+    timestamp         : float
+    """Monotonic capture time on the station clock, unit: [s]."""
+
+    joint_angle_q     : jt.Float32[np.ndarray, "n_arm_dof"]
+    """Joint positions, unit: [rad]."""
+    joint_angle_vel_dq: jt.Float32[np.ndarray, "n_arm_dof"]
+    """Joint velocities, unit: [rad/s]."""
+    joint_torque      : jt.Float32[np.ndarray, "n_arm_dof"]
+    """Joint torques, unit: [N*m]."""
+    joint_temperature : jt.Float32[np.ndarray, "n_arm_dof"]
+    """Motor temperatures, unit: [degC]."""
+    gripper_pos       : jt.Float32[np.ndarray, "n_gripper_dof"]
+    """Normalized gripper open fraction in [0, 1]; 0 = closed, 1 = open."""
+
+    def validate(self, spec: MotorizedArmSpec) -> None:
+        joint = jt.Float32[np.ndarray, f"{spec.n_arm_dof}"]
+        check_hint("joint_angle_q"     , self.joint_angle_q     , joint)
+        check_hint("joint_angle_vel_dq", self.joint_angle_vel_dq, joint)
+        check_hint("joint_torque"      , self.joint_torque      , joint)
+        check_hint("joint_temperature" , self.joint_temperature , joint)
+        check_hint("gripper_pos"       , self.gripper_pos       , jt.Float32[np.ndarray, f"{spec.n_gripper_dof}"])
+
+
+@jt.jaxtyped(typechecker=beartype)
+@dataclass(frozen=True)
+class MotorizedArm_JointPose:
+    """One absolute joint configuration of a single arm, e.g. the pose the arm walks to before an episode."""
+
+    joint_angle_q: jt.Float32[np.ndarray, "n_arm_dof"]
+    """Absolute joint positions, unit: [rad]."""
+
+    gripper_pos  : jt.Float32[np.ndarray, "n_gripper_dof"]
+    """Normalized gripper open fraction in [0, 1]; 0 = closed, 1 = open."""
+
+    def __post_init__(self) -> None:
+        if np.any(self.gripper_pos < 0.0) or np.any(self.gripper_pos > 1.0):
+            raise ValueError("gripper_pos must lie in [0, 1]")
+
+    def validate(self, spec: MotorizedArmSpec) -> None:
+        """Bind the dims to a concrete arm spec."""
+        check_hint("joint_angle_q", self.joint_angle_q, jt.Float32[np.ndarray, f"{spec.n_arm_dof}"])
+        check_hint("gripper_pos"  , self.gripper_pos  , jt.Float32[np.ndarray, f"{spec.n_gripper_dof}"])
+
+
+@jt.jaxtyped(typechecker=beartype)
+@dataclass(frozen=True)
+class MotorizedArm_ActionChunk:
+    """A horizon of absolute joint position targets for a single arm.
+
+    Construction enforces dtype / rank and a shared ``horizon`` between joints and gripper;
+    ``validate`` binds the dims to a concrete ``MotorizedArmSpec`` and horizon.
+
+    Row ``i`` is the target ``i / rate`` seconds after the chunk's reference timestamp
+    (both carried by the enclosing ``station.Station_ActionChunk``).
+    """
+
+    joint_angle_q: jt.Float32[np.ndarray, "horizon n_arm_dof"]
+    """Absolute joint position targets, unit: [rad]."""
+
+    gripper_pos  : jt.Float32[np.ndarray, "horizon n_gripper_dof"]
+    """Normalized gripper open fraction in [0, 1]; 0 = closed, 1 = open."""
+
+    def __post_init__(self) -> None:
+        if np.any(self.gripper_pos < 0.0) or np.any(self.gripper_pos > 1.0):
+            raise ValueError("gripper_pos must lie in [0, 1]")
+
+    @property
+    def horizon(self) -> int:
+        return int(self.joint_angle_q.shape[0])
+
+    def __len__(self) -> int:
+        return self.horizon
+
+    def __getitem__(self, index: int) -> MotorizedArm_JointPose:
+        """Row ``index`` of the chunk as a single joint pose. Negative indices count from the end."""
+        if not -self.horizon <= index < self.horizon:
+            raise IndexError(f"chunk index {index} out of range for horizon {self.horizon}")
+        return MotorizedArm_JointPose(self.joint_angle_q[index], self.gripper_pos[index])
+
+    def validate(self, spec: MotorizedArmSpec, horizon: int) -> None:
+        """Bind the dims to a concrete arm spec and horizon."""
+        check_hint("joint_angle_q", self.joint_angle_q, jt.Float32[np.ndarray, f"{horizon} {spec.n_arm_dof}"])
+        check_hint("gripper_pos"  , self.gripper_pos  , jt.Float32[np.ndarray, f"{horizon} {spec.n_gripper_dof}"])

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/camera.py

@@ -0,0 +1,65 @@
+import numpy as np
+import jaxtyping as jt
+from enum import Enum
+from dataclasses import dataclass
+from beartype import beartype
+
+from ..util.validation import check_hint
+
+
+class Camera_DeviceType(Enum):
+    Realsense_D405    = "realsense_d405"
+    """Realsense D405 Stereo Camera. Both imagers are RGB, so one observation carries 2 views (left, right)."""
+
+    @property
+    def n_views(self) -> int:
+        """Number of frames delivered per observation (1 for mono, 2 for a stereo pair)."""
+        return _CAMERA_N_VIEWS[self]
+
+
+_CAMERA_N_VIEWS: dict[Camera_DeviceType, int] = {
+    Camera_DeviceType.Realsense_D405: 2,
+}
+
+
+@dataclass(frozen=True)
+class Camera_DeviceSpec:
+    device_type     : Camera_DeviceType
+    """Device type of the camera"""
+
+    resolution      : tuple[int, int]
+    """Resolution of each view of the camera stream, in tuple of (height, width), unit: [px]"""
+
+    refresh_rate    : int
+    """Refresh rate of the camera, unit: [fps]"""
+
+    jpeg_compression: int | None
+    """jpeg compression quality value if specified, if None then no compression will be made."""
+
+    @property
+    def n_views(self) -> int:
+        return self.device_type.n_views
+
+
+@jt.jaxtyped(typechecker=beartype)
+@dataclass(frozen=True)
+class Camera_Observation:
+    """One decoded observation of a camera: ``b`` RGB views captured together (``b = 2`` for a stereo pair).
+
+    Frames are delivered already decoded regardless of ``jpeg_compression``. Construction enforces
+    ``uint8`` BHWC with 3 channels; ``validate`` binds (b, h, w) to the camera spec.
+    """
+
+    timestamp: float
+    """Monotonic capture time on the station clock, unit: [s]. Shared by all views."""
+
+    rgb      : jt.UInt8[np.ndarray, "b h w 3"]
+    """Views in device order (stereo: index 0 = left, 1 = right)."""
+
+    @property
+    def n_views(self) -> int:
+        return int(self.rgb.shape[0])
+
+    def validate(self, spec: Camera_DeviceSpec) -> None:
+        h, w = spec.resolution
+        check_hint("rgb", self.rgb, jt.UInt8[np.ndarray, f"{spec.n_views} {h} {w} 3"])

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/job.py

@@ -0,0 +1,22 @@
+"""What a customer submits to a station: the task, the policy, and how many
+episodes to run. The station executes the request as uploaded; it does not
+get to say what is being evaluated."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .task import Task_Spec
+
+from .policy import Policy
+
+
+@dataclass(frozen=True)
+class JobRequest:
+    task    : Task_Spec
+    policy  : Policy
+    episodes: int
+
+    def __post_init__(self) -> None:
+        if self.episodes <= 0:
+            raise ValueError("a job needs at least one episode")

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/policy.py

@@ -0,0 +1,32 @@
+"""The policy contract: what a caller hands to ``evaluate(station, task, policy)``.
+
+Lifecycle, driven by the station, once per episode:
+
+1. ``initialize(station, task)``: return the joint pose the arms walk to before the episode starts.
+   Load weights and read the task instruction here.
+2. ``infer(observation)``, repeated on every observation tick until the episode ends. Return a new
+   chunk, or ``None`` to keep executing the current one. A new chunk replaces the remainder of the
+   current one.
+3. ``reset()`` after the episode ends. Clear observation history and any unexecuted chunk so the
+   next ``initialize`` starts clean.
+
+Anything about how the policy runs (its rate, chunk length, replanning, image preprocessing) lives
+inside these methods and is not declared to the SDK.
+"""
+from abc import ABC, abstractmethod
+
+from .station import Station_Spec, Station_Observation, Station_ActionChunk, Station_JointPose
+from .task import Task_Spec
+
+
+class Policy(ABC):
+    @abstractmethod
+    def initialize(self, station: Station_Spec, task: Task_Spec) -> Station_JointPose:
+        """Prepare for one episode of ``task`` on ``station`` and return the pose the arms start from."""
+
+    @abstractmethod
+    def infer(self, observation: Station_Observation) -> Station_ActionChunk | None:
+        """A new action chunk for this tick, or ``None`` to continue executing the current chunk."""
+
+    def reset(self) -> None:
+        """Clear per-episode state. Called once after the episode ends. No-op by default."""

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/repository.py

@@ -0,0 +1,27 @@
+"""Registry of known station specs."""
+from . import arm, camera
+from .station import Bimanual_Station_Spec, Station_BaseSpec
+
+
+Bimanual_YAM_Station = Bimanual_Station_Spec(
+    base=Station_BaseSpec.Static,
+    arm =arm.MotorizedArmSpec(
+        arm         =arm.MotorizedArm_ArmType.YAM_6DoF,
+        gripper     =arm.MotorizedArm_GripperType.I2RT_linear_gripper_1DoF,
+        refresh_rate=100,
+        control_rate=100,
+    ),
+    wrist_camera=camera.Camera_DeviceSpec(
+        device_type     =camera.Camera_DeviceType.Realsense_D405,
+        resolution      =(360, 640), # (h, w); 2x downsample for 720p.
+        refresh_rate    =30,
+        jpeg_compression=80
+    ),
+    topdown_camera=camera.Camera_DeviceSpec(
+        device_type     =camera.Camera_DeviceType.Realsense_D405,
+        resolution      =(360, 640), # (h, w); 2x downsample for 720p.
+        refresh_rate    =30,
+        jpeg_compression=80
+    )
+)
+"""Static bimanual YAM station: two YAM 6-DoF arms with I2RT grippers, two wrist D405s, one top-down D405."""

se3labs_interface-0.0.1a1/src/se3labs/interface/eval/station.py

@@ -0,0 +1,179 @@
+"""Station specs (hardware capability) and the station-level observation / action containers.
+
+A station is a named collection of arms and cameras. Observations and actions are keyed by
+those names so that a bimanual and a single-arm station share one container type.
+
+The policy side of the contract is ``policy.Policy``: it consumes ``Station_Observation`` ticks and
+produces ``Station_ActionChunk``s. Anything about how the policy runs (its rate, chunk length,
+replanning, image preprocessing) is decided inside the policy, not declared to the SDK.
+"""
+from abc import ABC, abstractmethod
+from enum import Enum
+from dataclasses import dataclass
+
+from . import arm as arm_types
+from . import camera
+
+
+class Station_BaseSpec(Enum):
+    Static = "static"
+    """Static manipulation base (the arm is mounted on a fixed, non-mobile base)."""
+
+
+class Station_Spec(ABC):
+    base: Station_BaseSpec
+
+    @property
+    @abstractmethod
+    def arms(self) -> dict[str, arm_types.MotorizedArmSpec]:
+        """Arm name -> spec. Names are the keys of ``Station_Observation.arms`` and ``Station_ActionChunk.arms``."""
+
+    @property
+    @abstractmethod
+    def cameras(self) -> dict[str, camera.Camera_DeviceSpec]:
+        """Camera name -> spec. Names are the keys of ``Station_Observation.cameras``."""
+
+
+@dataclass(frozen=True)
+class Bimanual_Station_Spec(Station_Spec):
+    base          : Station_BaseSpec
+    
+    arm           : arm_types.MotorizedArmSpec
+    """Spec shared by the ``left`` and ``right`` arms."""
+    
+    wrist_camera  : camera.Camera_DeviceSpec
+    """Spec shared by the ``wrist_left`` and ``wrist_right`` cameras."""
+    
+    topdown_camera: camera.Camera_DeviceSpec
+
+    @property
+    def arms(self) -> dict[str, arm_types.MotorizedArmSpec]:
+        return {"left": self.arm, "right": self.arm}
+
+    @property
+    def cameras(self) -> dict[str, camera.Camera_DeviceSpec]:
+        return {
+            "wrist_left" : self.wrist_camera,
+            "wrist_right": self.wrist_camera,
+            "topdown"    : self.topdown_camera,
+        }
+
+
+@dataclass(frozen=True)
+class Single_Arm_Station_Spec(Station_Spec):
+    base          : Station_BaseSpec
+    arm           : arm_types.MotorizedArmSpec
+    wrist_camera  : camera.Camera_DeviceSpec
+
+    @property
+    def arms(self) -> dict[str, arm_types.MotorizedArmSpec]:
+        return {"arm": self.arm}
+
+    @property
+    def cameras(self) -> dict[str, camera.Camera_DeviceSpec]:
+        return {"wrist": self.wrist_camera}
+
+
+@dataclass(frozen=True)
+class Station_Observation:
+    """One policy-side observation: the latest tick of every arm and camera.
+
+    Arms and cameras run at different rates, so each entry carries its own timestamp; consumers that
+    need temporal alignment (world models, chunked policies) should use those rather than assume sync.
+    """
+
+    arms   : dict[str, arm_types.MotorizedArm_Observation]
+    cameras: dict[str, camera.Camera_Observation]
+
+    def validate(self, station: Station_Spec) -> None:
+        """Check every arm / camera of ``station`` is present with the shapes its spec promises."""
+        missing_arms = set(station.arms) - set(self.arms)
+        if missing_arms:
+            raise ValueError(f"observation missing arms: {sorted(missing_arms)}")
+        for name, spec in station.arms.items():
+            try:
+                self.arms[name].validate(spec)
+            except (TypeError, ValueError) as e:
+                raise type(e)(f"arm '{name}': {e}") from None
+
+        missing_cams = set(station.cameras) - set(self.cameras)
+        if missing_cams:
+            raise ValueError(f"observation missing cameras: {sorted(missing_cams)}")
+        for name, spec in station.cameras.items():
+            try:
+                self.cameras[name].validate(spec)
+            except (TypeError, ValueError) as e:
+                raise type(e)(f"camera '{name}': {e}") from None
+
+
+@dataclass(frozen=True)
+class Station_JointPose:
+    """One absolute joint configuration for every arm of a station."""
+
+    arms: dict[str, arm_types.MotorizedArm_JointPose]
+
+    def __post_init__(self) -> None:
+        if not self.arms:
+            raise ValueError("joint pose has no arms")
+
+    def validate(self, station: Station_Spec) -> None:
+        """Check the pose names exactly the station's arms with the right shapes."""
+        if set(self.arms) != set(station.arms):
+            raise ValueError(f"pose arms {sorted(self.arms)} != station arms {sorted(station.arms)}")
+        for name, spec in station.arms.items():
+            try:
+                self.arms[name].validate(spec)
+            except (TypeError, ValueError) as e:
+                raise type(e)(f"arm '{name}': {e}") from None
+
+
+@dataclass(frozen=True)
+class Station_ActionChunk:
+    """A chunk of ``horizon`` absolute joint-position targets for every arm of a station.
+
+    Row ``i`` of every arm is the target at ``timestamp + i / rate``. The station interpolates rows
+    up to each arm's ``control_rate`` and executes them until the policy returns a new chunk, which
+    replaces the remainder of the current one.
+    """
+
+    timestamp: float
+    """Station-clock time of the observation this chunk was predicted from, unit: [s]."""
+
+    rate     : float
+    """Spacing between consecutive rows, unit: [Hz]. Chosen by the policy per chunk."""
+
+    arms     : dict[str, arm_types.MotorizedArm_ActionChunk]
+
+    def __post_init__(self) -> None:
+        if self.rate <= 0:
+            raise ValueError(f"rate must be positive, got {self.rate}")
+        if not self.arms:
+            raise ValueError("action chunk has no arms")
+        horizons = {a.horizon for a in self.arms.values()}
+        if len(horizons) != 1:
+            raise ValueError(f"inconsistent horizons across arms: {horizons}")
+
+    @property
+    def horizon(self) -> int:
+        return next(iter(self.arms.values())).horizon
+
+    def __len__(self) -> int:
+        return self.horizon
+
+    def __getitem__(self, index: int) -> Station_JointPose:
+        """Row ``index`` of the chunk, for every arm, as one station pose. Negative indices count from the end."""
+        return Station_JointPose({name: chunk[index] for name, chunk in self.arms.items()})
+
+    def validate(self, station: Station_Spec) -> None:
+        """Check the chunk names exactly the station's arms, with the right shapes and a feasible rate."""
+        if set(self.arms) != set(station.arms):
+            raise ValueError(f"action arms {sorted(self.arms)} != station arms {sorted(station.arms)}")
+        h = self.horizon
+        for name, spec in station.arms.items():
+            if self.rate > spec.control_rate:
+                raise ValueError(f"arm '{name}': chunk rate {self.rate} Hz exceeds control_rate {spec.control_rate} Hz")
+            try:
+                self.arms[name].validate(spec, h)
+            except (TypeError, ValueError) as e:
+                raise type(e)(f"arm '{name}': {e}") from None
+