247 lines
8.5 KiB
Python
247 lines
8.5 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import h5py
|
|
import numpy as np
|
|
|
|
from xr_rm_teleop.single_arm_velocity_teleop import SingleArmVelocityTeleop
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EpisodeTrajectory:
|
|
source_path: Path
|
|
times_s: np.ndarray
|
|
target_poses: np.ndarray
|
|
initial_joints: np.ndarray
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class MethodSummary:
|
|
method: str
|
|
damping: float | None
|
|
success_rate: float
|
|
position_rmse_m: float
|
|
orientation_rmse_rad: float
|
|
max_joint_speed_deg_s: float
|
|
|
|
|
|
def _normalized_quaternion(values: np.ndarray) -> np.ndarray:
|
|
quaternion = np.asarray(values, dtype=float)
|
|
if quaternion.shape != (4,) or not np.isfinite(quaternion).all():
|
|
raise ValueError("quaternion must contain 4 finite values")
|
|
norm = float(np.linalg.norm(quaternion))
|
|
if norm <= 1e-12:
|
|
raise ValueError("quaternion norm must be positive")
|
|
return quaternion / norm
|
|
|
|
|
|
def _slerp_quaternion(
|
|
start: np.ndarray,
|
|
end: np.ndarray,
|
|
fraction: float,
|
|
) -> np.ndarray:
|
|
first = _normalized_quaternion(start)
|
|
second = _normalized_quaternion(end)
|
|
dot = float(np.dot(first, second))
|
|
if dot < 0.0:
|
|
second = -second
|
|
dot = -dot
|
|
dot = float(np.clip(dot, -1.0, 1.0))
|
|
if dot > 1.0 - 1e-8:
|
|
return _normalized_quaternion(
|
|
first + float(fraction) * (second - first)
|
|
)
|
|
angle = float(np.arccos(dot))
|
|
sine = float(np.sin(angle))
|
|
return _normalized_quaternion(
|
|
np.sin((1.0 - fraction) * angle) / sine * first
|
|
+ np.sin(fraction * angle) / sine * second
|
|
)
|
|
|
|
|
|
def resample_trajectory(
|
|
trajectory: EpisodeTrajectory,
|
|
sample_rate_hz: float,
|
|
) -> EpisodeTrajectory:
|
|
if not np.isfinite(sample_rate_hz) or sample_rate_hz <= 0.0:
|
|
raise ValueError("sample_rate_hz must be finite and positive")
|
|
source_times = np.asarray(trajectory.times_s, dtype=float)
|
|
poses = np.asarray(trajectory.target_poses, dtype=float)
|
|
if source_times.ndim != 1 or poses.shape != (source_times.size, 7):
|
|
raise ValueError("trajectory must contain N timestamps and N x 7 poses")
|
|
if source_times.size < 2 or np.any(np.diff(source_times) <= 0.0):
|
|
raise ValueError("trajectory timestamps must be strictly increasing")
|
|
|
|
duration = float(source_times[-1] - source_times[0])
|
|
count = int(round(duration * sample_rate_hz)) + 1
|
|
target_times = np.linspace(source_times[0], source_times[-1], count)
|
|
target_poses = np.empty((count, 7), dtype=float)
|
|
for axis in range(3):
|
|
target_poses[:, axis] = np.interp(
|
|
target_times,
|
|
source_times,
|
|
poses[:, axis],
|
|
)
|
|
for index, timestamp in enumerate(target_times):
|
|
right = int(np.searchsorted(source_times, timestamp, side="right"))
|
|
right = min(max(right, 1), source_times.size - 1)
|
|
left = right - 1
|
|
interval = source_times[right] - source_times[left]
|
|
fraction = float((timestamp - source_times[left]) / interval)
|
|
target_poses[index, 3:] = _slerp_quaternion(
|
|
poses[left, 3:], poses[right, 3:], fraction
|
|
)
|
|
target_poses[0] = poses[0]
|
|
target_poses[-1] = poses[-1]
|
|
return EpisodeTrajectory(
|
|
source_path=trajectory.source_path,
|
|
times_s=target_times - target_times[0],
|
|
target_poses=target_poses,
|
|
initial_joints=np.asarray(trajectory.initial_joints, dtype=float).copy(),
|
|
)
|
|
|
|
|
|
def _longest_true_run(mask: np.ndarray) -> slice:
|
|
values = np.asarray(mask, dtype=bool)
|
|
best_start = best_stop = start = 0
|
|
for index, enabled in enumerate(np.r_[values, False]):
|
|
if enabled:
|
|
continue
|
|
if index - start > best_stop - best_start:
|
|
best_start, best_stop = start, index
|
|
start = index + 1
|
|
if best_stop - best_start < 2:
|
|
raise ValueError("episode has no valid teleoperation run")
|
|
return slice(best_start, best_stop)
|
|
|
|
|
|
def load_episode(path: Path) -> EpisodeTrajectory:
|
|
source = Path(path).expanduser().resolve()
|
|
if not source.is_file() or source.suffix.lower() not in (".h5", ".hdf5"):
|
|
raise FileNotFoundError(f"episode not found: {source}")
|
|
with h5py.File(source, "r") as handle:
|
|
if str(handle.attrs.get("arm", "")) != "right_rm75":
|
|
raise ValueError("episode arm must be right_rm75")
|
|
if str(handle.attrs.get("pose_order", "")) != "x,y,z,qx,qy,qz,qw":
|
|
raise ValueError("episode pose_order is unsupported")
|
|
required = (
|
|
"debug/tcp/final_target_pose",
|
|
"debug/timestamps/control_monotonic_ns",
|
|
"debug/control/teleop_active",
|
|
"debug/control/action_valid",
|
|
"debug/control/command_sent",
|
|
"observations/qpos",
|
|
)
|
|
missing = [name for name in required if name not in handle]
|
|
if missing:
|
|
raise ValueError(f"episode datasets missing: {missing}")
|
|
poses = np.asarray(handle[required[0]], dtype=float)
|
|
timestamps = np.asarray(handle[required[1]], dtype=np.int64)
|
|
mask = np.logical_and.reduce(
|
|
[
|
|
np.asarray(handle[required[2]], dtype=bool),
|
|
np.asarray(handle[required[3]], dtype=bool),
|
|
np.asarray(handle[required[4]], dtype=bool),
|
|
]
|
|
)
|
|
qpos = np.asarray(handle[required[5]], dtype=float)
|
|
if poses.shape != (timestamps.size, 7) or qpos.shape != (timestamps.size, 8):
|
|
raise ValueError("episode arrays have inconsistent shapes")
|
|
if not np.isfinite(poses).all() or not np.isfinite(qpos).all():
|
|
raise ValueError("episode contains NaN/Inf")
|
|
selected = _longest_true_run(mask)
|
|
selected_times = timestamps[selected]
|
|
if np.any(np.diff(selected_times) <= 0):
|
|
raise ValueError("episode timestamps must be strictly increasing")
|
|
selected_poses = poses[selected].copy()
|
|
selected_poses[:, 3:] = np.asarray(
|
|
[_normalized_quaternion(value) for value in selected_poses[:, 3:]]
|
|
)
|
|
return EpisodeTrajectory(
|
|
source_path=source,
|
|
times_s=(selected_times - selected_times[0]) * 1e-9,
|
|
target_poses=selected_poses,
|
|
initial_joints=qpos[selected.start, :7].copy(),
|
|
)
|
|
|
|
|
|
def _rotation_z(angle: float) -> np.ndarray:
|
|
cosine, sine = math.cos(angle), math.sin(angle)
|
|
return np.asarray(
|
|
[
|
|
[cosine, -sine, 0.0],
|
|
[sine, cosine, 0.0],
|
|
[0.0, 0.0, 1.0],
|
|
]
|
|
)
|
|
|
|
|
|
def orientation_error_rad(actual: np.ndarray, target: np.ndarray) -> float:
|
|
delta = np.asarray(target) @ np.asarray(actual).T
|
|
cosine = float(np.clip((np.trace(delta) - 1.0) * 0.5, -1.0, 1.0))
|
|
return float(math.acos(cosine))
|
|
|
|
|
|
def normalized_joint_margin(
|
|
joints: np.ndarray,
|
|
lower: np.ndarray,
|
|
upper: np.ndarray,
|
|
) -> float:
|
|
values = np.asarray(joints, dtype=float)
|
|
lower_values = np.asarray(lower, dtype=float)
|
|
upper_values = np.asarray(upper, dtype=float)
|
|
span = upper_values - lower_values
|
|
if np.any(span <= 0.0):
|
|
raise ValueError("joint limits must have positive spans")
|
|
margins = np.minimum(
|
|
values - lower_values,
|
|
upper_values - values,
|
|
) / span
|
|
return float(np.min(margins))
|
|
|
|
|
|
def choose_dls_damping(candidates: list[MethodSummary]) -> float:
|
|
if not candidates or any(value.damping is None for value in candidates):
|
|
raise ValueError("DLS candidates must contain damping values")
|
|
selected = min(
|
|
candidates,
|
|
key=lambda value: (
|
|
-value.success_rate,
|
|
value.position_rmse_m / 0.002
|
|
+ value.orientation_rmse_rad / 0.005,
|
|
value.max_joint_speed_deg_s,
|
|
),
|
|
)
|
|
return float(selected.damping)
|
|
|
|
|
|
def limit_joint_command(
|
|
*,
|
|
target: np.ndarray,
|
|
previous_target: np.ndarray,
|
|
previous_velocity: np.ndarray,
|
|
max_speed: float,
|
|
max_acceleration: float,
|
|
dt: float,
|
|
) -> tuple[np.ndarray, np.ndarray, bool]:
|
|
limited_target, limited_velocity = (
|
|
SingleArmVelocityTeleop._limit_joint_command_step(
|
|
target=np.asarray(target, dtype=float).tolist(),
|
|
previous_target=np.asarray(previous_target, dtype=float).tolist(),
|
|
previous_velocity=np.asarray(previous_velocity, dtype=float).tolist(),
|
|
max_speed=max_speed,
|
|
max_acceleration=max_acceleration,
|
|
dt=dt,
|
|
)
|
|
)
|
|
target_array = np.asarray(limited_target, dtype=float)
|
|
velocity_array = np.asarray(limited_velocity, dtype=float)
|
|
return (
|
|
target_array,
|
|
velocity_array,
|
|
not np.allclose(target_array, target, atol=1e-12, rtol=0.0),
|
|
)
|