from __future__ import annotations from dataclasses import dataclass, replace from enum import Enum from typing import Dict, Mapping, Optional import numpy as np import pinocchio as pin from .kinematics import RM75Kinematics, validate_se3 from .robot_backend import RobotBackend from .solver import RM75IkSolver from .teleop_config import ArmName, ArmTeleopProfile from .types import IkOptions, teleop_joint_limits class SafetyState(str, Enum): IDLE = "idle" ACTIVE = "active" RESETTING = "resetting" FAULT = "fault" @dataclass(frozen=True) class ControllerSample: hand: ArmName grip: bool trigger: float position_m: np.ndarray quaternion_xyzw: np.ndarray def __post_init__(self) -> None: if self.hand not in ("left", "right"): raise ValueError("controller hand must be 'left' or 'right'") position = np.asarray(self.position_m, dtype=float).copy() quaternion = np.asarray(self.quaternion_xyzw, dtype=float).copy() if position.shape != (3,) or not np.all(np.isfinite(position)): raise ValueError("controller position must be finite with shape (3,)") if quaternion.shape != (4,) or not np.all(np.isfinite(quaternion)): raise ValueError("controller quaternion must be finite with shape (4,)") norm = float(np.linalg.norm(quaternion)) if norm <= 1e-9: raise ValueError("controller quaternion has zero norm") quaternion /= norm if not np.isfinite(self.trigger): raise ValueError("controller trigger must be finite") position.setflags(write=False) quaternion.setflags(write=False) object.__setattr__(self, "position_m", position) object.__setattr__(self, "quaternion_xyzw", quaternion) @classmethod def from_message( cls, message, expected_arm: Optional[ArmName] = None ) -> "ControllerSample": message_hand = str(getattr(message, "hand", "")).strip().lower() hand = expected_arm or message_hand if expected_arm is not None and message_hand and message_hand != expected_arm: raise ValueError( f"controller message hand {message_hand!r} does not match {expected_arm!r}" ) pose = message.pose return cls( hand=hand, grip=bool(message.grip), trigger=float(message.trigger), position_m=np.array( [pose.position.x, pose.position.y, pose.position.z], dtype=float ), quaternion_xyzw=np.array( [ pose.orientation.x, pose.orientation.y, pose.orientation.z, pose.orientation.w, ], dtype=float, ), ) @dataclass(frozen=True) class MappedTarget: target_tcp: pin.SE3 clamped: bool just_engaged: bool @dataclass(frozen=True) class ControlCycleResult: state: SafetyState commanded_arms: tuple[ArmName, ...] = () targets_tcp: Optional[Mapping[ArmName, pin.SE3]] = None reason: str = "" class RelativePoseMapper: def __init__(self, profile: ArmTeleopProfile) -> None: self.profile = profile self.active = False self._controller_start_position: Optional[np.ndarray] = None self._controller_start_rotation: Optional[np.ndarray] = None self._robot_start_tcp: Optional[pin.SE3] = None self._filtered_position: Optional[np.ndarray] = None self._filtered_rotation: Optional[np.ndarray] = None self._last_position: Optional[np.ndarray] = None self._last_rotation: Optional[np.ndarray] = None def reset(self) -> None: self.active = False self._controller_start_position = None self._controller_start_rotation = None self._robot_start_tcp = None self._filtered_position = None self._filtered_rotation = None self._last_position = None self._last_rotation = None @staticmethod def _rotation(sample: ControllerSample) -> np.ndarray: x, y, z, w = sample.quaternion_xyzw return pin.Quaternion(w, x, y, z).matrix() def map( self, sample: ControllerSample, current_tcp: pin.SE3, dt: float, ) -> MappedTarget: if sample.hand != self.profile.arm: raise ValueError("controller sample was routed to the wrong arm") validate_se3(current_tcp, "current_tcp") if not np.isfinite(dt) or dt <= 0.0: raise ValueError("control dt must be finite and positive") rotation = self._rotation(sample) if not self.active: self.active = True self._controller_start_position = sample.position_m.copy() self._controller_start_rotation = rotation.copy() self._robot_start_tcp = current_tcp.copy() self._filtered_position = current_tcp.translation.copy() self._filtered_rotation = current_tcp.rotation.copy() self._last_position = current_tcp.translation.copy() self._last_rotation = current_tcp.rotation.copy() return MappedTarget(current_tcp.copy(), False, True) assert self._controller_start_position is not None assert self._controller_start_rotation is not None assert self._robot_start_tcp is not None delta = sample.position_m - self._controller_start_position mapped_delta = self.profile.xr_to_robot @ delta raw_position = ( self._robot_start_tcp.translation + self.profile.scale * mapped_delta ) raw_position = np.where( np.asarray(self.profile.enable_position_axes, dtype=bool), raw_position, self._robot_start_tcp.translation, ) position, clamped = self._clamp_workspace(raw_position) position = self._filter_position(position) position, limited = self._limit_position_step(position, dt) position, final_clamped = self._clamp_workspace(position) target_rotation = self._target_rotation(rotation) target_rotation = self._filter_rotation(target_rotation) target_rotation, rotation_limited = self._limit_rotation_step( target_rotation, dt ) self._last_position = position.copy() self._last_rotation = target_rotation.copy() return MappedTarget( pin.SE3(target_rotation, position), clamped or limited or final_clamped or rotation_limited, False, ) def _clamp_workspace(self, target: np.ndarray) -> tuple[np.ndarray, bool]: result = np.clip( np.asarray(target, dtype=float), self.profile.workspace_min, self.profile.workspace_max, ) min_radius, max_radius = self.profile.cylinder_radius_limit if result[2] < self.profile.low_z_threshold: min_radius = max(min_radius, self.profile.low_z_min_radius) radius = float(np.hypot(result[0], result[1])) if radius > max_radius: result[:2] *= max_radius / radius elif radius < min_radius: if radius > 1e-9: result[:2] *= min_radius / radius else: result[:2] = [min_radius, 0.0] changed = not np.allclose(result, target, atol=1e-12, rtol=0.0) return result, changed def _filter_position(self, target: np.ndarray) -> np.ndarray: assert self._filtered_position is not None assert self._last_position is not None if np.linalg.norm(target - self._last_position) < self.profile.deadband_m: target = self._last_position distance = float(np.linalg.norm(target - self._filtered_position)) ratio = min( 1.0, distance / self.profile.target_filter_fast_threshold_m ) alpha = self.profile.target_filter_alpha + ratio * ( self.profile.target_filter_alpha_fast - self.profile.target_filter_alpha ) self._filtered_position = ( alpha * target + (1.0 - alpha) * self._filtered_position ) return self._filtered_position.copy() def _limit_position_step( self, target: np.ndarray, dt: float ) -> tuple[np.ndarray, bool]: assert self._last_position is not None delta = target - self._last_position distance = float(np.linalg.norm(delta)) maximum = self.profile.max_linear_speed_m_s * dt if distance <= maximum or distance <= 1e-12: return target, False return self._last_position + delta * (maximum / distance), True def _target_rotation(self, controller_rotation: np.ndarray) -> np.ndarray: assert self._controller_start_rotation is not None assert self._robot_start_tcp is not None if not self.profile.enable_orientation_control: return self._robot_start_tcp.rotation.copy() xr_delta = controller_rotation @ self._controller_start_rotation.T matrix = self.profile.xr_to_robot robot_delta = matrix @ xr_delta @ matrix.T target = robot_delta @ self._robot_start_tcp.rotation axes = np.asarray(self.profile.enable_orientation_axes, dtype=bool) if not np.all(axes): target_rpy = pin.rpy.matrixToRpy(target) start_rpy = pin.rpy.matrixToRpy(self._robot_start_tcp.rotation) target = pin.rpy.rpyToMatrix(*np.where(axes, target_rpy, start_rpy)) return target def _filter_rotation(self, target: np.ndarray) -> np.ndarray: assert self._filtered_rotation is not None assert self._last_rotation is not None if ( np.linalg.norm(pin.log3(self._last_rotation.T @ target)) < self.profile.orientation_deadband_rad ): target = self._last_rotation delta = pin.log3(self._filtered_rotation.T @ target) self._filtered_rotation = self._filtered_rotation @ pin.exp3( self.profile.orientation_filter_alpha * delta ) return self._filtered_rotation.copy() def _limit_rotation_step( self, target: np.ndarray, dt: float ) -> tuple[np.ndarray, bool]: assert self._last_rotation is not None delta = pin.log3(self._last_rotation.T @ target) angle = float(np.linalg.norm(delta)) maximum = self.profile.max_orientation_speed_rad_s * dt if angle <= maximum or angle <= 1e-12: return target, False return self._last_rotation @ pin.exp3(delta * (maximum / angle)), True class DualArmQpTeleopController: def __init__( self, robot: RobotBackend, profiles: Mapping[ArmName, ArmTeleopProfile], control_rate_hz: float = 90.0, ik_options: Optional[IkOptions] = None, ) -> None: if set(profiles) != {"left", "right"}: raise ValueError("profiles must contain left and right arms") if not np.isfinite(control_rate_hz) or control_rate_hz <= 0.0: raise ValueError("control_rate_hz must be finite and positive") self.robot = robot self.profiles = dict(profiles) self.dt = 1.0 / control_rate_hz self.ik_options = ik_options or IkOptions( max_iterations=120, time_limit_sec=0.008, ) self.kinematics = { arm: RM75Kinematics(limits=teleop_joint_limits()) for arm in ("left", "right") } self.solvers = { arm: RM75IkSolver(self.kinematics[arm]) for arm in ("left", "right") } self.mappers = { arm: RelativePoseMapper(self.profiles[arm]) for arm in ("left", "right") } self._latest: Dict[ArmName, tuple[ControllerSample, float]] = {} self._fault_reason = "" self._reset_waiting_for_release = False self._closed = False self.robot.connect() @property def safety_state(self) -> SafetyState: if self._fault_reason: return SafetyState.FAULT if self._reset_waiting_for_release: return SafetyState.RESETTING if any(mapper.active for mapper in self.mappers.values()): return SafetyState.ACTIVE return SafetyState.IDLE def update_controller( self, message, timestamp_sec: float, expected_arm: Optional[ArmName] = None, ) -> None: self.update_sample( ControllerSample.from_message(message, expected_arm), timestamp_sec ) def update_sample(self, sample: ControllerSample, timestamp_sec: float) -> None: if not np.isfinite(timestamp_sec): raise ValueError("controller timestamp must be finite") self._latest[sample.hand] = (sample, float(timestamp_sec)) def reject_input(self, reason: str) -> None: self._trip_fault(f"invalid controller input: {reason}") def reset_to_initial(self) -> ControlCycleResult: """Reset the backend and require fresh dual-grip release before rearming.""" if self._closed: raise RuntimeError("controller is closed") for mapper in self.mappers.values(): mapper.reset() try: self.robot.stop(("left", "right")) self.robot.reset_to_initial() except Exception as exc: self._trip_fault(f"reset/backend failure: {exc}") raise RuntimeError(f"failed to reset robot backend: {exc}") from exc self._latest.clear() self._fault_reason = "" self._reset_waiting_for_release = True return ControlCycleResult( SafetyState.RESETTING, reason="reset complete; waiting for fresh dual-grip release", ) def step(self, timestamp_sec: float) -> ControlCycleResult: if self._closed: raise RuntimeError("controller is closed") if not np.isfinite(timestamp_sec): return self._trip_fault("control timestamp is non-finite") if self._fault_reason: if self._can_clear_fault(timestamp_sec): self._fault_reason = "" return ControlCycleResult(SafetyState.IDLE, reason="fault cleared") return ControlCycleResult(SafetyState.FAULT, reason=self._fault_reason) if self._reset_waiting_for_release: if self._can_finish_reset(timestamp_sec): self._reset_waiting_for_release = False self._latest.clear() return ControlCycleResult(SafetyState.IDLE, reason="reset rearmed") return ControlCycleResult( SafetyState.RESETTING, reason="waiting for fresh dual-grip release", ) try: state = self.robot.read_joint_positions() commands: Dict[ArmName, np.ndarray] = {} targets: Dict[ArmName, pin.SE3] = {} for arm in ("left", "right"): latest = self._latest.get(arm) mapper = self.mappers[arm] if latest is None: continue sample, sample_time = latest age = timestamp_sec - sample_time if age < -1e-6: return self._trip_fault(f"{arm} controller timestamp is in the future") if not sample.grip: if mapper.active: mapper.reset() self.robot.stop((arm,)) continue if age > self.profiles[arm].command_timeout_sec: return self._trip_fault(f"{arm} controller input timed out") q_current = state.positions_rad[arm] current_tcp = self.kinematics[arm].forward( q_current, self.profiles[arm].tool_from_flange ) mapped = mapper.map(sample, current_tcp, self.dt) flange_target = ( mapped.target_tcp * self.profiles[arm].tool_from_flange.inverse() ) arm_ik_options = replace( self.ik_options, joint_limit_mid_weight=self.profiles[arm].qp_w_limit_mid, joint_motion_weights=tuple( float(value) for value in self.profiles[arm].qp_joint_motion_weights ), joint_step_limits_rad=tuple( float(value) for value in self.profiles[arm].qp_joint_step_limits_rad ), ) result = self.solvers[arm].solve( flange_target, q_current, arm_ik_options, ) if not result.success or result.q is None: return self._trip_fault( f"{arm} IK failed: {result.status.value}: {result.message}" ) max_step = self.profiles[arm].joint_max_speed_rad_s * self.dt q_command = np.clip(result.q, q_current - max_step, q_current + max_step) limits = self.kinematics[arm].limits q_command = np.clip(q_command, limits.lower, limits.upper) commands[arm] = q_command targets[arm] = mapped.target_tcp if commands: self.robot.command_joint_positions(commands) for arm, target in targets.items(): self.robot.set_target_tcp_pose(arm, target) return ControlCycleResult( self.safety_state, tuple(commands), targets or None, ) except Exception as exc: return self._trip_fault(f"control/backend failure: {exc}") def _can_clear_fault(self, timestamp_sec: float) -> bool: if set(self._latest) != {"left", "right"}: return False for arm, (sample, sample_time) in self._latest.items(): if sample.grip: return False if timestamp_sec - sample_time > self.profiles[arm].command_timeout_sec: return False for mapper in self.mappers.values(): mapper.reset() return True def _can_finish_reset(self, timestamp_sec: float) -> bool: if set(self._latest) != {"left", "right"}: return False for arm, (sample, sample_time) in self._latest.items(): if sample.grip: return False age = timestamp_sec - sample_time if age < -1e-6 or age > self.profiles[arm].command_timeout_sec: return False for mapper in self.mappers.values(): mapper.reset() return True def _trip_fault(self, reason: str) -> ControlCycleResult: self._fault_reason = str(reason) self._reset_waiting_for_release = False for mapper in self.mappers.values(): mapper.reset() try: self.robot.stop(("left", "right")) except Exception: pass return ControlCycleResult(SafetyState.FAULT, reason=self._fault_reason) def stop(self) -> None: self._reset_waiting_for_release = False for mapper in self.mappers.values(): mapper.reset() self.robot.stop(("left", "right")) def close(self) -> None: if self._closed: return try: self.stop() finally: self.robot.close() self._closed = True