Implement dual-arm teleoperation with RM75 QP controller and MuJoCo backend
This commit is contained in:
429
ik_qp/src/rm75_ik/teleop_control.py
Normal file
429
ik_qp/src/rm75_ik/teleop_control.py
Normal file
@ -0,0 +1,429 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
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"
|
||||
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._closed = False
|
||||
self.robot.connect()
|
||||
|
||||
@property
|
||||
def safety_state(self) -> SafetyState:
|
||||
if self._fault_reason:
|
||||
return SafetyState.FAULT
|
||||
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 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)
|
||||
|
||||
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()
|
||||
)
|
||||
result = self.solvers[arm].solve(
|
||||
flange_target, q_current, self.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 _trip_fault(self, reason: str) -> ControlCycleResult:
|
||||
self._fault_reason = str(reason)
|
||||
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:
|
||||
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
|
||||
Reference in New Issue
Block a user