493 lines
19 KiB
Python
493 lines
19 KiB
Python
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
|