diff --git a/src/lerobot_robot_ufactory/scripts/uf_lerobot_eval.py b/src/lerobot_robot_ufactory/scripts/uf_lerobot_eval.py index ee8772c..f698f20 100644 --- a/src/lerobot_robot_ufactory/scripts/uf_lerobot_eval.py +++ b/src/lerobot_robot_ufactory/scripts/uf_lerobot_eval.py @@ -38,6 +38,7 @@ from lerobot.configs.policies import PreTrainedConfig from lerobot.scripts.lerobot_record import DatasetRecordConfig from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot_robot_ufactory.utils.utils import init_keyboard_listener +from lerobot_robot_ufactory.utils.action_safety import ActionSafetyConfig, ActionSafetyGuard from lerobot_robot_ufactory.devices.umi.vive_tracker.transformations import Transformations @@ -120,9 +121,11 @@ class EvalConfig: return ["policy"] -def eval_loop(cfg: EvalConfig, relative=False, rx_continuous=False): +def eval_loop(cfg: EvalConfig, relative=False, rx_continuous=False, safety_guard: ActionSafetyGuard | None = None): init_logging() logging.info(pformat(asdict(cfg))) + if safety_guard is not None: + safety_guard.log_config() robot = make_robot_from_config(cfg.robot) @@ -261,6 +264,8 @@ def eval_loop(cfg: EvalConfig, relative=False, rx_continuous=False): prev_robot_dict[key] = {'type': 1 if is_tcp else 0, 'pose': np.array(pose)} prev_action_dict[key] = {'type': 1 if is_tcp else 0, 'pose': np.array(pose)} + safety_halted = False + while True: start_loop_t = time.perf_counter() @@ -269,6 +274,12 @@ def eval_loop(cfg: EvalConfig, relative=False, rx_continuous=False): print("\n********** Policy Eval Episode (Reset) **********") break + # Safety halt: stop inference and stop sending actions, + # wait for the operator to press the right arrow key to resume + if safety_halted: + precise_sleep(sleep_time_s) + continue + # Get robot observation obs = robot.get_observation() @@ -398,6 +409,16 @@ def eval_loop(cfg: EvalConfig, relative=False, rx_continuous=False): # gripper_raw = (future_gripper_norm + 1) / 2 * (_gripper_max - _gripper_min) + _gripper_min # robot_action_to_send['right.gripper.pos'] = 1.0 if gripper_raw > 0.4 else 0.0 + # Safety check: any violation triggers an e-stop; the action is + # NOT sent and the loop waits for the operator to press right arrow + if safety_guard is not None: + violation = safety_guard.check(robot_action_to_send, curr_robot_dict, keys) + if violation is not None: + safety_halted = True + logging.error(f"*** SAFETY HALT *** {violation}") + print(f"\n*** SAFETY HALT *** {violation}\nAction was NOT sent. Press right arrow (->) to reset and resume, ESC to exit.") + continue + robot.send_action(robot_action_to_send) dt_s = time.perf_counter() - start_loop_t @@ -418,11 +439,15 @@ def main(): parser = argparse.ArgumentParser(description='configuration args') parser.add_argument('--relative', action='store_true', help='is relative motion or not') parser.add_argument('--rx_continuous', action='store_true', help='rx continuous or not') + parser.add_argument('--enable_safety', action='store_true', help='enable action safety guard (thresholds in ActionSafetyConfig)') args, unknown = parser.parse_known_args() sys.argv = [sys.argv[0]] + unknown register_third_party_plugins() cfg = get_cfg() - eval_loop(cfg, args.relative, args.rx_continuous) + # Action safety guard: tune thresholds in ActionSafetyConfig directly. + # Any violation triggers an e-stop; press right arrow to resume. + safety_guard = ActionSafetyGuard(ActionSafetyConfig(enabled=args.enable_safety)) + eval_loop(cfg, args.relative, args.rx_continuous, safety_guard) if __name__ == "__main__": diff --git a/src/lerobot_robot_ufactory/utils/action_safety.py b/src/lerobot_robot_ufactory/utils/action_safety.py new file mode 100644 index 0000000..7fe6b80 --- /dev/null +++ b/src/lerobot_robot_ufactory/utils/action_safety.py @@ -0,0 +1,130 @@ +"""Model-agnostic action safety checks. + +Validates the final action dict right before it is sent to the robot; any +failed check triggers an emergency halt. Only relies on the action key +naming convention: `[prefix.]pose.x/y/z/rx/ry/rz`, `*.gripper.pos`, so it +works with any policy (ACT / DP / pi0 / custom backbones) and with both +single-arm and multi-arm setups. +""" + +import logging +import math +from dataclasses import dataclass + +import numpy as np + +from lerobot_robot_ufactory.devices.umi.vive_tracker.transformations import Transformations + +POSE_AXES = ("x", "y", "z", "rx", "ry", "rz") + + +@dataclass +class ActionSafetyConfig: + """Configuration for ActionSafetyGuard. + + All limits are compared against the robot's current actual pose per + control step. `workspace_min`/`workspace_max` use the same unit as the + pose commands (e.g. mm) and must be provided together. + """ + + enabled: bool = True + max_step_mm: float = 25.0 + max_rot_step: float = 0.35 + workspace_min: list | None = None + workspace_max: list | None = None + + def __post_init__(self): + if (self.workspace_min is None) != (self.workspace_max is None): + raise ValueError("workspace_min and workspace_max must be provided together") + + +def _rotation_delta_norm(rot_prev, rot_curr) -> float: + """Relative rotation angle (rad) between two axis-angle rotations. + + Uses rotation matrices for the diff to avoid the ±π discontinuity + of subtracting raw rotvecs. + """ + R_prev = Transformations.rxryrz_to_rotation_matrix(*rot_prev) + R_curr = Transformations.rxryrz_to_rotation_matrix(*rot_curr) + R_delta = R_prev.T @ R_curr + delta = Transformations.rotation_matrix_to_rxryrz(R_delta) + return float(np.linalg.norm(delta)) + + +class ActionSafetyGuard: + """Safety checks for the action dict about to be sent to the robot. + + check() returns None when the action is safe; otherwise it returns a + human-readable violation reason (the caller should trigger an e-stop). + """ + + def __init__(self, config: ActionSafetyConfig | None = None): + config = config or ActionSafetyConfig() + self.config = config + self.enabled = config.enabled + self.max_step_mm = config.max_step_mm + self.max_rot_step = config.max_rot_step + self.workspace_min = None if config.workspace_min is None else np.asarray(config.workspace_min, dtype=np.float64) + self.workspace_max = None if config.workspace_max is None else np.asarray(config.workspace_max, dtype=np.float64) + + def check(self, action: dict, curr_robot_dict: dict, keys) -> str | None: + """Check the action dict. `action` holds absolute pose commands and + `curr_robot_dict` holds each arm's current actual pose.""" + if not self.enabled: + return None + + for key in keys: + prefix = f".{key}" if key else "" + pose_keys = [f"{prefix}pose.{axis}" for axis in POSE_AXES] + if not all(k in action for k in pose_keys): + continue # not TCP-pose controlled (e.g. joint-only arm), skip pose checks + + target = np.array([action[k] for k in pose_keys], dtype=np.float64) + + # 1. Numeric validity: NaN / Inf + if not np.all(np.isfinite(target)): + return f"[{key or 'arm'}] action contains NaN/Inf: {target.tolist()}" + + # 2. Single-step delta limit, measured from the robot's current actual pose + curr = curr_robot_dict.get(key) + if curr is not None and curr.get("type") == 1: + curr_pose = np.asarray(curr["pose"], dtype=np.float64) + pos_dist = float(np.linalg.norm(target[:3] - curr_pose[:3])) + if pos_dist > self.max_step_mm: + return ( + f"[{key or 'arm'}] position step {pos_dist:.1f}mm exceeds limit " + f"{self.max_step_mm}mm (current {curr_pose[:3].tolist()} -> target {target[:3].tolist()})" + ) + rot_delta = _rotation_delta_norm(curr_pose[3:6], target[3:6]) + if rot_delta > self.max_rot_step: + return ( + f"[{key or 'arm'}] rotation step {rot_delta:.3f}rad exceeds limit " + f"{self.max_rot_step}rad" + ) + + # 3. Workspace bounding box (optional) + if self.workspace_min is not None: + pos = target[:3] + if np.any(pos < self.workspace_min) or np.any(pos > self.workspace_max): + return ( + f"[{key or 'arm'}] target position {pos.tolist()} outside workspace " + f"[{self.workspace_min.tolist()}, {self.workspace_max.tolist()}]" + ) + + # Gripper numeric check + for k, v in action.items(): + if "gripper" in k and isinstance(v, (int, float)) and not math.isfinite(v): + return f"[{k}] gripper action contains NaN/Inf: {v}" + + return None + + def log_config(self): + ws = ( + f"[{self.workspace_min.tolist()}, {self.workspace_max.tolist()}]" + if self.workspace_min is not None + else "not set" + ) + logging.info( + f"ActionSafetyGuard: max_step={self.max_step_mm}mm, " + f"max_rot_step={self.max_rot_step}rad, workspace={ws}" + )