client/jam_gello_client.py

#!/usr/bin/env python3
"""ONE FILE. Copy it to the GELLO robot PC and run it. Nothing else is installed.

    # on the GPU box (already verified there):
    #   tools/serve_resident.py ... --profile foundation --port 8000
    # then forward the port to this PC:
    #   ssh -L 8000:127.0.0.1:8000 <gpu-host>

    python jam_gello_client.py --selftest                     # no robot, no camera, no server
    python jam_gello_client.py --task 1 --mode dry-run --record runs
    python jam_gello_client.py --task 1 --mode step    --record runs --workspace ...
    python jam_gello_client.py --task 1 --mode run     --record runs --workspace ...

Everything `jamkit` and `gello_bridge` provide is inlined below, so this file has no
JAM dependency at all. It needs, on this PC: numpy, pyrealsense2, requests (or nothing --
urllib is used), and the `gello` package that already drives the arm.

WHAT IT DOES
    RealSense 317422074704 -> pad_to_5x4 -> PNG -> POST /act
    measured joints        -> Panda FK at panda_link8 -> pose10 -> POST /act
    32 absolute flange targets come back at 15 Hz
    -> safety filter -> damped-least-squares IK -> command_joint_state([q1..q7, aperture])

WHAT IT DOES NOT DO
    It does not replace the hardware E-stop, joint limits, collision detection or force
    thresholds. Those stay in your control stack. Keep a hand on the E-stop.

THE ADAPTERS HAVE NEVER BEEN VALIDATED ON A ROBOT. Run --mode dry-run first, then step.
"""
import argparse
import base64
import datetime
import io
import json
import sys
import time
import urllib.request

import numpy as np

# --------------------------------------------------------------------------- #
# profile constants -- must match the training data exactly
# --------------------------------------------------------------------------- #
DT = 1.0 / 15.0                     # the model's targets are 15 Hz apart
HEAD_SERIAL = '317422074704'        # third-person camera. The wrist camera (821312060498)
                                    # is NOT part of single_external_no_wrist_v1.
GRIPPER_MAX_WIDTH = 0.08

# Flange (panda_link8) -> Franka Hand TCP: 0.1034 m along flange z, rotated -45 deg about z.
# JAM emits FLANGE targets; command_delta_pose works in the TCP frame that obs eef_pos/eef_quat
# report, so --control delta has to convert. Confirm against the F_T_EE configured on YOUR
# robot (Desk -> Settings -> End-effector) before trusting it.
_c, _s45 = np.cos(-np.pi / 4), np.sin(-np.pi / 4)
FRANKA_HAND_F_T_EE = np.array([[_c, -_s45, 0, 0], [_s45, _c, 0, 0], [0, 0, 1, 0.1034], [0, 0, 0, 1]])
# TWO DIFFERENT GRIPPER ENCODINGS. They are not interchangeable.
#
# command_joint_state  -> send the 0..1 APERTURE straight through (1 = open).
#   gello/robots/panda.py does its own binarisation:
#       gripper_action = -1 if gripper_action >= 0.8 else 1
#   and per TRAINING.md -1 = open (Move), +1 = close (Grasp). So >=0.8 means OPEN.
#   This is the identical threshold that binarised the human's GELLO trigger while
#   the demonstrations were recorded, and the training action channel is already
#   binary {0.0, 1.0}. The model regresses continuous values: apply hysteresis
#   first, then send a clean 0/1 aperture through this API.
#   Sending +-1 here is inverted: -1 ("open") fails >=0.8 and becomes a CLOSE.
#
# command_delta_pose   -> send +-1, +1 = close, -1 = open.
#   That is the pi0.5 eval script's convention (gripper_action = 1.0 if width <= 0.07).
GRIPPER_OPEN, GRIPPER_CLOSE = -1.0, +1.0     # the DELTA-path flag only
PANDA_OPEN_AT = 0.8                          # panda.py's own threshold, for reference

# Where rollouts are written: <RECORD_ROOT>/task_<N>/session_<8 hex>/.
# THIS LINE IS THE SETTING. Override per-run with --record.
# It is RELATIVE, so it follows your working directory: run the script from the
# directory you keep it in and rollouts land beside it. Set an absolute path here,
# or pass --record, if you launch it from somewhere else.
RECORD_ROOT = 'runs'

# Legacy median from 13 raw recording starts. This is NOT guaranteed to match
# the 100 cleaned training episode starts: their median distances are 5.5-8.2 cm.
# Keep the existing reset here; validate a task-specific start pose on site.
# Flange lands at [0.596, 0.055, 0.461].  NOTE this is NOT the reset pose in
# run_gello_env_eval_pi05_ee.py, whose flange is [0.499, -0.004, 0.632] -- about
# 20 cm away, mostly in z.  That pose was chosen for a different policy.
HOME_JOINTS = np.array([-0.1223, 0.1887, 0.2364, -1.7202, -0.0613, 1.7389, 0.9986])

# Bounding box of the FLANGE over all 13 recorded sessions (min [0.457,-0.220,0.211],
# max [0.746,0.356,0.548]), plus a 5 cm margin. The demonstrations are the best evidence
# of where this arm can safely go on this rig, so this is the default rather than
# something the operator has to measure. Override with --workspace after measuring, and
# note the z floor: the demos genuinely reach down to 0.211 m.
DEMO_WORKSPACE = [0.407, -0.270, 0.161, 0.796, 0.406, 0.598]

PROMPTS = {
    1: 'stack the green block on the orange block, then stack the yellow block on top',
    2: 'straighten the yellow chain',
    3: 'pick up the scoop and pour its contents into the green bowl',
    4: 'pick up the bread and put it in the red bowl',
    5: 'pick up the green pepper and put it in the white bowl',
}

# Craig MODIFIED DH (a, d, alpha), then the flange offset. Identical to
# jam.corpus.kinematics.PANDA_MDH, which produced the training poses.
# The standard-DH form looks similar and is silently wrong by over a metre.
PANDA_MDH = ((0.0, 0.333, 0.0), (0.0, 0.0, -np.pi / 2), (0.0, 0.316, np.pi / 2),
             (0.0825, 0.0, np.pi / 2), (-0.0825, 0.384, -np.pi / 2),
             (0.0, 0.0, np.pi / 2), (0.088, 0.0, np.pi / 2))
PANDA_FLANGE = 0.107
JOINT_LIMITS = np.array([[-2.8973, 2.8973], [-1.7628, 1.7628], [-2.8973, 2.8973],
                         [-3.0718, -0.0698], [-2.8973, 2.8973], [-0.0175, 3.7525],
                         [-2.8973, 2.8973]])


# --------------------------------------------------------------------------- #
# rotation helpers (inlined from jamkit.rot)
# --------------------------------------------------------------------------- #
def matrix_to_quat_xyzw(R):
    R = np.asarray(R, np.float64)
    t = np.trace(R)
    if t > 0:
        s = np.sqrt(t + 1.0) * 2
        q = np.array([(R[2, 1] - R[1, 2]) / s, (R[0, 2] - R[2, 0]) / s, (R[1, 0] - R[0, 1]) / s, 0.25 * s])
    else:
        i = int(np.argmax(np.diag(R)))
        j, k = (i + 1) % 3, (i + 2) % 3
        s = np.sqrt(R[i, i] - R[j, j] - R[k, k] + 1.0) * 2
        q = np.zeros(4)
        q[i], q[j], q[k] = 0.25 * s, (R[j, i] + R[i, j]) / s, (R[k, i] + R[i, k]) / s
        q[3] = (R[k, j] - R[j, k]) / s
    return q / np.linalg.norm(q)


def quat_xyzw_to_matrix(q):
    q = np.asarray(q, np.float64)
    n = np.linalg.norm(q)
    if q.shape != (4,) or not np.isfinite(q).all() or abs(n - 1) > 1e-3:
        raise ValueError('Expected a unit xyzw quaternion')
    x, y, z, w = q / n
    return np.array([
        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]])


def rot6d_to_matrix(r6):
    """Gram-Schmidt on [column0, column1] -- COLUMNS, not rows."""
    r6 = np.asarray(r6, np.float64)
    a, b = r6[..., :3], r6[..., 3:6]
    na = np.linalg.norm(a, axis=-1, keepdims=True)
    if np.any(na < 1e-6):
        raise ValueError('Degenerate first rotation column')
    x = a / na
    b = b - x * (x * b).sum(-1, keepdims=True)
    nb = np.linalg.norm(b, axis=-1, keepdims=True)
    if np.any(nb < 1e-6):
        raise ValueError('Degenerate second rotation column')
    y = b / nb
    return np.stack([x, y, np.cross(x, y)], axis=-1)


def matrix_to_rot6d(R):
    R = np.asarray(R, np.float64)
    return np.concatenate([R[..., :, 0], R[..., :, 1]], axis=-1)


def geodesic_deg(Ra, Rb):
    Ra, Rb = np.asarray(Ra, np.float64), np.asarray(Rb, np.float64)
    c = (np.einsum('...ij,...ij->...', Ra, Rb) - 1.0) / 2.0
    return np.degrees(np.arccos(np.clip(c, -1.0, 1.0)))


def axis_angle_to_matrix(v):
    v = np.asarray(v, np.float64)
    th = np.linalg.norm(v)
    if th < 1e-12:
        return np.eye(3)
    k = v / th
    K = np.array([[0, -k[2], k[1]], [k[2], 0, -k[0]], [-k[1], k[0], 0]])
    return np.eye(3) + np.sin(th) * K + (1 - np.cos(th)) * (K @ K)


def matrix_to_axis_angle(R):
    q = matrix_to_quat_xyzw(R)
    if q[3] < 0:
        q = -q
    s = np.linalg.norm(q[:3])
    if s < 1e-12:
        return np.zeros(3)
    return q[:3] / s * (2 * np.arctan2(s, q[3]))


def limit_rotation(R_from, R_to, max_deg):
    v = matrix_to_axis_angle(np.asarray(R_from).T @ np.asarray(R_to))
    th = np.linalg.norm(v)
    lim = np.radians(max_deg)
    if th <= lim:
        return np.asarray(R_to, np.float64), False
    return np.asarray(R_from) @ axis_angle_to_matrix(v / th * lim), True


# --------------------------------------------------------------------------- #
# pose10 (inlined from jamkit.pose)
# --------------------------------------------------------------------------- #
def pose10_from_matrix(T_O_F, aperture):
    T = np.asarray(T_O_F, np.float64)
    if T.shape != (4, 4) or not np.isfinite(T).all() or not 0 <= aperture <= 1:
        raise ValueError('Expected a finite 4x4 transform and aperture in [0,1]')
    R = T[:3, :3]
    if np.abs(R.T @ R - np.eye(3)).max() > 1e-3 or np.linalg.det(R) < 0:
        raise ValueError('Rotation block is not a proper rotation')
    return np.r_[T[:3, 3], matrix_to_rot6d(R), aperture].astype(np.float32)


def unpack_pose10(p):
    p = np.asarray(p, np.float64)
    if p.shape != (10,) or not np.isfinite(p).all():
        raise ValueError('Expected a finite pose10')
    R = rot6d_to_matrix(p[3:9])
    return p[:3].copy(), R, matrix_to_quat_xyzw(R), float(p[9])


def check_pose10(x, tol=0.02):
    x = np.asarray(x, np.float32)
    if x.ndim < 1 or x.shape[-1] != 10 or not np.isfinite(x).all():
        raise ValueError('Expected finite pose10 rows')
    u, v = x[..., 3:6], x[..., 6:9]
    if (np.abs(np.linalg.norm(u, axis=-1) - 1) > tol).any() or (np.abs(np.linalg.norm(v, axis=-1) - 1) > tol).any() \
            or (np.abs((u * v).sum(-1)) > tol).any():
        raise ValueError('Rotation columns must be orthonormal (first column, then second column)')
    if ((x[..., 9] < 0) | (x[..., 9] > 1)).any():
        raise ValueError('Aperture must lie in [0,1] (0 closed, 1 open)')
    return x


# --------------------------------------------------------------------------- #
# image (inlined from jamkit.image)
# --------------------------------------------------------------------------- #
def pad_to_5x4(rgb):
    """Letterbox to 5:4 WITHOUT resampling. The server then stretches to 320x256, so
    padding first makes that stretch an aspect-preserving resize. Mandatory for
    the foundation profile; a mismatch here is silent."""
    rgb = np.asarray(rgb)
    if rgb.dtype != np.uint8 or rgb.ndim != 3 or rgb.shape[-1] != 3:
        raise ValueError('Expected an HxWx3 uint8 RGB frame')
    h, w = rgb.shape[:2]
    if w * 4 >= h * 5:
        H = int(np.ceil(w * 4 / 5))
        out = np.zeros((H, w, 3), np.uint8)
        top = (H - h) // 2
        out[top:top + h] = rgb
    else:
        W = int(np.ceil(h * 5 / 4))
        out = np.zeros((h, W, 3), np.uint8)
        left = (W - w) // 2
        out[:, left:left + w] = rgb
    return out


def bgr_to_rgb(frame):
    """RealSense/OpenCV deliver BGR. The checkpoint was trained on RGB; a swap is silent."""
    frame = np.asarray(frame)
    if frame.ndim != 3 or frame.shape[-1] != 3:
        raise ValueError('Expected HxWx3')
    return np.ascontiguousarray(frame[..., ::-1])


def encode_png(rgb):
    """PNG, never JPEG: JPEG artefacts are something the model never saw."""
    try:
        from PIL import Image
        buf = io.BytesIO()
        Image.fromarray(rgb).save(buf, format='PNG')
        return buf.getvalue()
    except ImportError:
        import cv2
        ok, enc = cv2.imencode('.png', rgb[..., ::-1])   # cv2 wants BGR
        if not ok:
            raise RuntimeError('PNG encoding failed')
        return enc.tobytes()


# --------------------------------------------------------------------------- #
# kinematics (inlined from gello_bridge)
# --------------------------------------------------------------------------- #
def _link(a, d, alpha, theta):
    ct, st, ca, sa = np.cos(theta), np.sin(theta), np.cos(alpha), np.sin(alpha)
    return np.array([[ct, -st, 0.0, a],
                     [st * ca, ct * ca, -sa, -d * sa],
                     [st * sa, ct * sa, ca, d * ca],
                     [0.0, 0.0, 0.0, 1.0]])


def panda_fk(q):
    """Base->FLANGE (panda_link8) transform. NOT the hand TCP, which is 0.1034 m / -45 deg away."""
    q = np.asarray(q, np.float64).reshape(-1)
    if q.shape != (7,) or not np.isfinite(q).all():
        raise ValueError('Expected seven finite joint angles')
    T = np.eye(4)
    for j, (a, d, alpha) in enumerate(PANDA_MDH):
        T = T @ _link(a, d, alpha, q[j])
    return T @ _link(0.0, PANDA_FLANGE, 0.0, 0.0)


def panda_jacobian(q):
    q = np.asarray(q, np.float64).reshape(-1)
    if q.shape != (7,):
        raise ValueError('Expected seven joint angles')
    T = np.eye(4)
    origins, axes = [], []
    for j, (a, d, alpha) in enumerate(PANDA_MDH):
        T = T @ _link(a, d, alpha, q[j])
        origins.append(T[:3, 3].copy())
        axes.append(T[:3, 2].copy())
    p_e = (T @ _link(0.0, PANDA_FLANGE, 0.0, 0.0))[:3, 3]
    J = np.zeros((6, 7))
    for j in range(7):
        J[:3, j] = np.cross(axes[j], p_e - origins[j])
        J[3:, j] = axes[j]
    return J


def solve_ik(target_pose10, q_seed, *, tol_m=0.002, tol_deg=0.6, damping=0.05,
             max_iters=60, step_clip=0.2):
    """Damped least squares to a FLANGE pose -> (q[7], residual_m, residual_deg).

    Raises rather than returning a half-solved pose: that would be a wrong command,
    not an approximate one."""
    target = np.asarray(target_pose10, np.float64).reshape(-1)
    R_t, p_t = rot6d_to_matrix(target[3:9]), target[:3]
    q = np.clip(np.asarray(q_seed, np.float64).reshape(-1)[:7], JOINT_LIMITS[:, 0], JOINT_LIMITS[:, 1])
    for _ in range(max_iters):
        T = panda_fk(q)
        e = np.concatenate([p_t - T[:3, 3], matrix_to_axis_angle(T[:3, :3].T @ R_t)])
        e[3:] = T[:3, :3] @ e[3:]
        if np.linalg.norm(e[:3]) <= tol_m and geodesic_deg(T[:3, :3], R_t) <= tol_deg:
            return q, float(np.linalg.norm(e[:3])), float(geodesic_deg(T[:3, :3], R_t))
        J = panda_jacobian(q)
        dq = J.T @ np.linalg.solve(J @ J.T + damping ** 2 * np.eye(6), e)
        n = np.linalg.norm(dq)
        if n > step_clip:
            dq *= step_clip / n
        q = np.clip(q + dq, JOINT_LIMITS[:, 0], JOINT_LIMITS[:, 1])
    T = panda_fk(q)
    raise RuntimeError('IK did not converge: %.4f m / %.2f deg from the target'
                       % (np.linalg.norm(p_t - T[:3, 3]), geodesic_deg(T[:3, :3], R_t)))


def state_pose10(obs):
    """Observation -> pose10, built exactly as the training data was: forward kinematics
    of the MEASURED joints at panda_link8. Recorded O_T_EE is deliberately not used so
    that deployment and training share one code path."""
    if 'joint_positions' not in obs:
        raise ValueError("observation has no 'joint_positions'")
    jp = np.asarray(obs['joint_positions'], np.float64).reshape(-1)
    if jp.shape[0] < 8 or not np.isfinite(jp[:8]).all():
        raise ValueError('Expected eight finite values: seven joints and a gripper fraction')
    aperture = float(np.clip(jp[7], 0.0, 1.0))    # already width/0.08 in this stack
    return check_pose10(pose10_from_matrix(panda_fk(jp[:7]), aperture))


# --------------------------------------------------------------------------- #
# TCP frame + delta commands -- matching example_real_world_scripts/client_side/
#   run_gello_env_eval_pi05_ee.py (get_ee_states / abs_to_local_delta)
# --------------------------------------------------------------------------- #
def tcp_from_flange(position, R_flange):
    """Flange pose -> Franka Hand TCP pose, the frame obs eef_pos/eef_quat reports."""
    T = np.eye(4)
    T[:3, :3], T[:3, 3] = R_flange, position
    T = T @ FRANKA_HAND_F_T_EE
    return T[:3, 3].copy(), T[:3, :3].copy()


def tcp_from_obs(obs):
    """Measured TCP pose, read exactly as the pi0.5 eval script's get_ee_states does."""
    if 'eef_pos' in obs and 'eef_quat' in obs:
        return (np.asarray(obs['eef_pos'], np.float64).reshape(3),
                quat_xyzw_to_matrix(np.asarray(obs['eef_quat'], np.float64).reshape(4)))
    if 'ee_pos_quat' in obs:
        v = np.asarray(obs['ee_pos_quat'], np.float64).reshape(-1)
        return v[:3].copy(), quat_xyzw_to_matrix(v[3:7])
    raise ValueError("observation has no 'eef_pos'/'eef_quat' or 'ee_pos_quat'")


def matrix_to_rot6d_rows(R):
    """First two ROWS, flattened -- dp_gs.dataset.utils.rot_mat_to_rot_6d (rot_mat[:, :2, :]).

    THIS IS NOT JAM's CONVENTION.  JAM pose10 carries the first two COLUMNS; the lab's
    own `command_delta_pose` carries the first two ROWS, and its inverse rebuilds the
    matrix as rows 0,1 plus cross(row0, row1).  Mixing them is a silent ~90 deg error
    that raises nothing.  Use matrix_to_rot6d() for anything going to the POLICY, and
    this one for anything going to the ROBOT's delta command.
    """
    R = np.asarray(R, np.float64)
    return np.concatenate([R[0, :], R[1, :]])


def abs_to_local_delta(target_pos, target_R, current_pos, current_R, gripper_flag):
    """Absolute TCP target -> the 10-vector command_delta_pose expects.

    [delta_pos(3) in the CURRENT TCP frame, delta_rot6d(6, ROW convention), gripper(1)]
    -- the same conversion as abs_to_local_delta() in the pi0.5 eval script
    (current_rot.inv().apply(dp) and current_rot.inv() * target_rot, then
    dp_gs.rot_mat_to_rot_6d), with the gripper flag supplied by our hysteresis instead
    of that script's raw-width threshold.
    """
    delta_pos = current_R.T @ (np.asarray(target_pos, np.float64) - np.asarray(current_pos, np.float64))
    delta_rot6d = matrix_to_rot6d_rows(current_R.T @ target_R)
    return np.concatenate([delta_pos, delta_rot6d, [float(gripper_flag)]])


# --------------------------------------------------------------------------- #
# safety (inlined from jamkit.safety)
# --------------------------------------------------------------------------- #
class SafetyFilter:
    def __init__(self, lo, hi, max_step_m=0.015, max_step_deg=4.0, max_jump_m=0.10, max_jump_deg=45.0):
        self.lo, self.hi = np.asarray(lo, np.float64), np.asarray(hi, np.float64)
        if self.lo.shape != (3,) or self.hi.shape != (3,) or (self.hi <= self.lo).any():
            raise ValueError('workspace bounds must be 3-vectors with hi > lo')
        self.max_step_m, self.max_step_deg = float(max_step_m), float(max_step_deg)
        self.max_jump_m, self.max_jump_deg = float(max_jump_m), float(max_jump_deg)

    def check_chunk(self, chunk, current):
        """Reject a whole chunk whose FIRST target is implausibly far from the robot.
        That almost always means the frame or the frame convention is wrong."""
        chunk, current = np.asarray(chunk, np.float64), np.asarray(current, np.float64)
        if chunk.ndim != 2 or chunk.shape[1] != 10 or not np.isfinite(chunk).all():
            return False, 'non-finite or malformed chunk'
        d = float(np.linalg.norm(chunk[0, :3] - current[:3]))
        a = float(geodesic_deg(rot6d_to_matrix(current[3:9]), rot6d_to_matrix(chunk[0, 3:9])))
        if d > self.max_jump_m or a > self.max_jump_deg:
            return False, 'first target is %.3f m / %.1f deg from the current pose' % (d, a)
        return True, ''

    def step(self, target, current):
        target, current = np.asarray(target, np.float64), np.asarray(current, np.float64)
        flags = {'clamped_workspace': False, 'limited_translation': False, 'limited_rotation': False}
        p = np.clip(target[:3], self.lo, self.hi)
        flags['clamped_workspace'] = bool((p != target[:3]).any())
        d = p - current[:3]
        n = np.linalg.norm(d)
        if n > self.max_step_m:
            p = current[:3] + d / n * self.max_step_m
            flags['limited_translation'] = True
        R, clipped = limit_rotation(rot6d_to_matrix(current[3:9]), rot6d_to_matrix(target[3:9]), self.max_step_deg)
        flags['limited_rotation'] = bool(clipped)
        out = np.r_[p, matrix_to_rot6d(R), np.clip(target[9], 0.0, 1.0)]
        return out.astype(np.float32), flags


class GripperHysteresis:
    """Debounce continuous policy aperture into a stable desired open/close state.
    update() reports edges for logging. Joint commands still carry the desired state
    every tick; the robot backend must avoid reissuing blocking Grasp on every tick.
    """

    def __init__(self, open_above=0.6, close_below=0.4, initially_open=True):
        if not 0 <= close_below < open_above <= 1:
            raise ValueError('need 0 <= close_below < open_above <= 1')
        self.open_above, self.close_below, self.is_open = open_above, close_below, bool(initially_open)

    @property
    def flag(self):
        return GRIPPER_OPEN if self.is_open else GRIPPER_CLOSE

    def update(self, aperture):
        a = float(aperture)
        if self.is_open and a < self.close_below:
            self.is_open = False
            return 'close'
        if not self.is_open and a > self.open_above:
            self.is_open = True
            return 'open'
        return None


# --------------------------------------------------------------------------- #
# hardware
# --------------------------------------------------------------------------- #
class HeadCamera:
    """The third-person RealSense, opened BY SERIAL so a replugged USB cannot silently
    swap it for the wrist camera."""

    def __init__(self, serial=HEAD_SERIAL, width=640, height=480, fps=30, warmup=60):
        import pyrealsense2 as rs
        self._rs = rs
        available = {d.get_info(rs.camera_info.serial_number) for d in rs.context().devices}
        if serial not in available:
            raise RuntimeError('RealSense %s not found. Connected: %s' % (serial, sorted(available) or 'none'))
        self.pipeline = rs.pipeline()
        cfg = rs.config()
        cfg.enable_device(serial)
        # rgb8, read straight through, NO channel swap -- byte-for-byte the path
        # experiments/run_gello_env.py used to RECORD the demonstrations:
        #     config1.enable_stream(rs.stream.color, 640, 480, rs.format.rgb8, 30)
        #     obs["camera1_rgb"] = np.asanyarray(color_frame1.get_data()).copy()
        # Requesting bgr8 and swapping is arithmetically the same but is a second place
        # to get the channel order wrong, and a swap here is silent.
        cfg.enable_stream(rs.stream.color, width, height, rs.format.rgb8, fps)
        self.pipeline.start(cfg)
        # 60 frames (~2 s) is MEASURED, not guessed. The D435I's auto-white-balance
        # converges slowly; on this rig the white table reads, against the demos'
        # neutral 163.8/163.2/159.7:
        #     2 frames  G-R +44.4     30 frames  G-R +13.1
        #    10 frames  G-R +36.5     60 frames  G-R  +0.0   <- matches the demos
        # A short warm-up gives every frame a green cast, which is silent and which
        # the colour-keyed prompts ("the green pepper", "the white bowl") care about.
        for _ in range(warmup):
            self.pipeline.wait_for_frames()

    def get_rgb(self):
        frames = self.pipeline.wait_for_frames()
        color = frames.get_color_frame()
        if not color:
            raise RuntimeError('no colour frame from the head camera')
        return np.asanyarray(color.get_data()).copy()

    def close(self):
        try:
            self.pipeline.stop()
        except Exception:
            pass


class Arm:
    """The GELLO ZMQ robot, with both command paths the lab already uses.

    `command_joint_state` is the BASE ZMQClientRobot call and the one GELLO
    teleoperation itself issues every tick (TRAINING.md: "T_t command_joint_state(...)
    -> send action to deoxys").  It is therefore the best-tested path on this rig, and
    for JAM it is also the exact inverse of how the training targets were built --
    `action` is FK of the COMMANDED joints, so IK recovers the quantity that was
    recorded.

    `command_delta_pose` is NOT part of the base robot client: the pi0.5 eval scripts
    add it in ZMQOpenPIClientRobot because their policy emits end-effector poses.  It is
    reproduced here verbatim, but it needs the robot node to handle that message type.
    """

    def __init__(self, host='127.0.0.1', port=6001):
        from gello.zmq_core.robot_node import ZMQClientRobot
        self.robot = ZMQClientRobot(port=port, host=host)
        self.last_aperture = None

    def observation(self):
        return self.robot.get_observations()

    def send_joints(self, q7, aperture):
        """Eighth element is the 0..1 APERTURE, exactly as GELLO teleoperation sent it.
        panda.py binarises it at >=0.8. Do NOT pass a +-1 flag here -- it inverts."""
        a = float(np.clip(aperture, 0.0, 1.0))
        self.robot.command_joint_state(np.concatenate([q7, [a]]).astype(np.float64))
        self.last_aperture = a

    def send_delta(self, delta10):
        """Verbatim the ZMQOpenPIClientRobot.command_delta_pose envelope."""
        import pickle
        delta10 = np.asarray(delta10, np.float64).reshape(-1)
        if delta10.shape != (10,) or not np.isfinite(delta10).all():
            raise ValueError('delta must be 10 finite numbers')
        msg = {'type': 'command_delta_pose', 'delta_action': delta10.tolist()}
        self.robot._socket.send(pickle.dumps(msg))
        resp = pickle.loads(self.robot._socket.recv())
        if not resp.get('success', False):
            raise RuntimeError('command_delta_pose failed: %s' % resp.get('error', 'unknown'))
        self.last_aperture = 1.0 if delta10[9] < 0 else 0.0

    def stop(self):
        """Hold position: re-command the joints the arm is already at."""
        if self.last_aperture is None:
            return  # This instance has issued no motion command.
        try:
            jp = np.asarray(self.observation()['joint_positions'], np.float64).reshape(-1)
            # Width while holding an object is not the desired open/close command.
            self.robot.command_joint_state(np.r_[jp[:7], self.last_aperture])
        except Exception as e:
            print('stop(): could not re-command current joints:', e)


# --------------------------------------------------------------------------- #
# recording -- same layout as the GELLO demonstrations
# --------------------------------------------------------------------------- #
class Recorder:
    """Writes the rollout in the SAME layout as the demo recordings.

        <root>/task_<N>/session_<8 hex>/
            data00000000.h5        every key the robot reports, plus policy/* extras
            serial_<serial>.mp4    the head camera, full resolution

    Matching the demo layout means a rollout goes through the existing tooling
    unchanged and a good one can be folded straight back into the training set.

    The policy extras are the things the demos have no column for, and they are the
    reason there is no separate log file:

        policy/chunk         (T, 32, 10)  the full predicted chunk on replan rows
        policy/replanned     (T, 1)       1 where a fresh inference happened
        policy/infer_s       (T, 1)       server latency on those rows, else nan
        policy/accepted      (T, 1)       1 if the chunk passed check_chunk
        policy/sent_pose10   (T, 10)      the filtered target actually commanded
        policy/raw_pose10    (T, 10)      the same target BEFORE the safety filter
        policy/ik_residual   (T, 2)       metres, degrees
        policy/flags         (T, 3)       clamped_workspace, limited_translation, limited_rotation
        policy/t_unix        (T, 1)       real timestamp; the mp4 rate is nominal
    """

    KEYS = ('chunk', 'replanned', 'infer_s', 'accepted', 'sent_pose10', 'raw_pose10',
            'ik_residual', 'flags', 't_unix', 'chunk_index', 'command_s', 'loop_s',
            'gripper_command', 'gripper_measured', 'position_error_m')

    def __init__(self, root, task, serial, fps=15.0, attrs=None):
        import os
        import uuid
        # <root>/task_N/session_<8 hex>/ -- the demos' own layout AND their own naming
        # (session_0dea3fb8, session_70ff1960, ...), so a rollout drops in beside the
        # demonstrations it is meant to be compared with. Random rather than a timestamp
        # because two runs started in the same second must not collide; the wall-clock
        # time is kept in the h5 attrs and in policy/t_unix, so nothing is lost.
        for _ in range(100):
            d = os.path.join(root, task, 'session_' + uuid.uuid4().hex[:8])
            if not os.path.exists(d):
                break
        else:
            raise RuntimeError('could not find a free session name under %s' % root)
        self.dir = d
        os.makedirs(self.dir)
        self.serial, self.fps, self.attrs = serial, float(fps), dict(attrs or {})
        self.writer = None
        self.obs_rows, self.pol_rows = {}, {k: [] for k in self.KEYS}
        print('recording to %s' % self.dir)

    def write(self, frame, obs, **kw):
        import cv2
        if self.writer is None:
            h, w = frame.shape[:2]
            self.writer = cv2.VideoWriter(self.dir + '/serial_%s.mp4' % self.serial,
                                          cv2.VideoWriter_fourcc(*'mp4v'), self.fps, (w, h))
        self.writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
        for k, v in obs.items():
            try:
                arr = np.asarray(v, np.float64).reshape(-1)
            except (ValueError, TypeError):
                continue                     # non-numeric keys (e.g. the demos' time string)
            self.obs_rows.setdefault(k, []).append(arr)
        def opt(key, default):
            # kw.get(key, default) is wrong here: callers pass an explicit None for
            # "no value this row", which would slip past the default.
            v = kw.get(key)
            return default if v is None else v

        nan10 = np.full(10, np.nan)
        row = {'chunk': np.asarray(opt('chunk', np.full((32, 10), np.nan)), np.float64),
               'replanned': np.array([1.0 if kw.get('replanned') else 0.0]),
               'infer_s': np.array([float(opt('infer_s', np.nan))]),
               'accepted': np.array([1.0 if opt('accepted', True) else 0.0]),
               'sent_pose10': np.asarray(opt('sent_pose10', nan10), np.float64),
               'raw_pose10': np.asarray(opt('raw_pose10', nan10), np.float64),
               'ik_residual': np.asarray(opt('ik_residual', [np.nan, np.nan]), np.float64),
               'flags': np.asarray(opt('flags', [0, 0, 0]), np.float64),
               't_unix': np.array([time.time()])}
        for key in ('chunk_index', 'command_s', 'loop_s', 'gripper_command',
                    'gripper_measured', 'position_error_m'):
            row[key] = np.array([float(opt(key, np.nan))])
        for k in self.KEYS:
            self.pol_rows[k].append(row[k])

    def close(self):
        if self.writer is not None:
            self.writer.release()
        n = len(self.pol_rows['t_unix'])
        if not n:
            print('nothing recorded')
            return
        import h5py
        path = self.dir + '/data00000000.h5'
        with h5py.File(path, 'w') as f:
            for k, rows in self.obs_rows.items():
                if len({len(r) for r in rows}) == 1 and len(rows) == n:
                    f.create_dataset(k, data=np.stack(rows))
            for k in self.KEYS:
                f.create_dataset('policy/' + k, data=np.stack(self.pol_rows[k]))
            for k, v in self.attrs.items():
                f.attrs[k] = v
            f.attrs['camera_head_serial'] = self.serial
            f.attrs['control_hz'] = self.fps
            f.attrs['recorded_at'] = datetime.datetime.now().isoformat(timespec='seconds')
        print('wrote %d rows to %s' % (n, path))


def reset_to_home(arm, home=None, aperture=1.0, max_step_rad=0.01, dwell=0.01, limit=300):
    """Interpolate slowly to the legacy HOME pose, gripper open.

    Same shape as the reset in run_gello_env_eval_pi05_ee.py -- linear in joint
    space, <=0.01 rad per command, 10 ms apart -- but to OUR home pose, and with
    the gripper as a 0..1 aperture so the single sweep also opens the hand.
    (That script needed a separate open command precisely because it sent -1,
    which command_joint_state reads as a close.)
    """
    home = HOME_JOINTS if home is None else np.asarray(home, np.float64).reshape(7)
    cur = np.asarray(arm.observation()['joint_positions'], np.float64).reshape(-1)[:7]
    gap = float(np.abs(cur - home).max())
    n = int(np.clip(round(gap / max_step_rad), 2, limit))
    print('reset: %.3f rad from home, %d steps' % (gap, n))
    for q in np.linspace(cur, home, n):
        arm.send_joints(q, aperture)
        time.sleep(dwell)
    settled = np.asarray(arm.observation()['joint_positions'], np.float64).reshape(-1)[:7]
    err = float(np.abs(settled - home).max())
    print('reset: settled %.4f rad from home' % err)
    return err


# --------------------------------------------------------------------------- #
# policy service
# --------------------------------------------------------------------------- #
def query(url, rgb, state, prompt, seed, timeout=120):
    body = json.dumps({'rgb_base64': base64.b64encode(encode_png(rgb)).decode(),
                       'state': [float(v) for v in state],
                       'prompt': prompt, 'seed': int(seed)}).encode()
    req = urllib.request.Request(url, data=body, headers={'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        out = json.loads(r.read())
    if 'error' in out:
        raise RuntimeError('policy server: ' + out['error'])
    actions = np.asarray(out['actions'], np.float32)
    if actions.shape != (32, 10) or not np.isfinite(actions).all():
        raise ValueError('Expected exactly 32 finite pose10 targets')
    return actions


# --------------------------------------------------------------------------- #
# control loop
# --------------------------------------------------------------------------- #
def run(a):
    prompt = a.prompt if a.prompt else PROMPTS[a.task]
    ws = a.workspace if a.workspace else DEMO_WORKSPACE
    print('prompt:    %r' % prompt)
    print('mode:      %s' % a.mode)
    print('workspace: %s%s' % ([round(v, 3) for v in ws],
                               '' if a.workspace else '   (default: demo bounding box + 5 cm)'))
    print('budget:    %d executed steps' % a.max_steps)
    print('control:   %s' % ('command_delta_pose (TCP frame)' if a.control == 'delta'
                             else 'command_joint_state (IK from flange)'))

    safety = SafetyFilter(ws[:3], ws[3:], a.max_step_m, a.max_step_deg)
    grip = GripperHysteresis(a.gripper_open_above, a.gripper_close_below, initially_open=True)
    cam = HeadCamera(a.camera_serial)
    arm = Arm(a.robot_host, a.robot_port)
    task_dir = 'task_%d' % a.task if a.task else 'task_custom'
    rec = None if a.no_record else Recorder(a.record, task_dir, a.camera_serial, 1.0 / DT,
                                            {'prompt': prompt, 'mode': a.mode,
                                             'execute_steps': a.execute_steps, 'url': a.url,
                                             'control': a.control, 'seed': a.seed,
                                             'gripper_open_above': a.gripper_open_above,
                                             'gripper_close_below': a.gripper_close_below})
    if a.mode != 'dry-run' and not a.no_reset:
        # Never in dry-run: that mode must not move the robot at all.
        reset_to_home(arm)
    if a.mode != 'dry-run':
        grip.is_open = bool(state_pose10(arm.observation())[9] >= PANDA_OPEN_AT)
    steps = 0          # EXECUTED targets -- the budget the operator actually cares about
    chunk_id = 0
    try:
        while steps < a.max_steps:
            obs = arm.observation()
            state = state_pose10(obs)
            frame = cam.get_rgb()

            t0 = time.monotonic()
            chunk = query(a.url, pad_to_5x4(frame), state, prompt, seed=a.seed, timeout=a.timeout)
            latency = time.monotonic() - t0
            ok, reason = safety.check_chunk(chunk, state)
            print('chunk %3d  latency %5.2fs  first target %.3f m away  accepted=%s %s'
                  % (chunk_id, latency, float(np.linalg.norm(chunk[0, :3] - state[:3])), ok, reason))
            chunk_id += 1

            if a.mode == 'dry-run':
                # No motion. The row still carries the frame, the state and the whole
                # predicted chunk, which is exactly what you inspect afterwards.
                if rec:
                    rec.write(frame, obs, chunk=chunk, replanned=True, infer_s=latency, accepted=ok)
                steps += 1
                if not ok:
                    print('REJECTED -- stopping. Usually a frame-convention error.')
                    break
                time.sleep(1.0)
                continue

            if not ok:
                if rec:
                    rec.write(frame, obs, chunk=chunk, replanned=True, infer_s=latency, accepted=False)
                print('REJECTED -- stopping. Usually a frame-convention error.')
                break

            n_exec = min(1 if a.mode == 'step' else a.execute_steps, a.max_steps - steps)
            chunk_t0 = time.monotonic()
            for k in range(n_exec):
                deadline = chunk_t0 + k * DT
                now = time.monotonic()
                if now - deadline > DT:
                    # A blocking hand/robot/camera call invalidated this chunk's timing.
                    # Re-observe and replan; never jump across an unexecuted close edge.
                    print('CONTROL_OVERRUN: %.3fs late at target %d; replan without skipping targets'
                          % (now - deadline, k))
                    break
                if now < deadline:
                    time.sleep(deadline - now)
                tick_t0 = time.monotonic()
                obs = arm.observation()
                current = state_pose10(obs)
                q_seed = np.asarray(obs['joint_positions'], np.float64).reshape(-1)[:7]
                target, flags = safety.step(chunk[k], current)
                res_m = res_deg = np.nan
                command = grip.update(float(target[9]))
                command_t0 = time.monotonic()
                if a.control == 'joint':
                    # The regressor is continuous even though training labels are binary.
                    # Apply the same hysteresis on both paths; encode its state as 0/1
                    # for panda.py (NOT the delta API's opposite-polarity +/-1 flag).
                    try:
                        q7, res_m, res_deg = solve_ik(target, q_seed)
                    except RuntimeError as e:
                        print('IK failed at k=%d: %s -- stopping.' % (k, e))
                        return
                    arm.send_joints(q7, 1.0 if grip.is_open else 0.0)
                else:
                    # The path the pi0.5 evaluations used on this rig: a local-frame
                    # delta in the TCP frame, tracked by the robot's own controller.
                    # This API DOES want the +-1 flag, so the hysteresis applies here.
                    tgt_pos, tgt_R = tcp_from_flange(target[:3], rot6d_to_matrix(target[3:9]))
                    cur_pos, cur_R = tcp_from_obs(obs)
                    arm.send_delta(abs_to_local_delta(tgt_pos, tgt_R, cur_pos, cur_R, grip.flag))
                command_s = time.monotonic() - command_t0
                if command:
                    print('  gripper -> %s (command %.3fs)' % (command, command_s))
                if rec:
                    rec.write(cam.get_rgb() if k else frame, obs,
                              chunk=chunk if k == 0 else None,
                              replanned=(k == 0), infer_s=latency if k == 0 else np.nan,
                              accepted=True, sent_pose10=target, raw_pose10=chunk[k],
                              chunk_index=k, command_s=command_s,
                              loop_s=time.monotonic()-tick_t0,
                              gripper_command=1.0 if grip.is_open else 0.0,
                              gripper_measured=float(current[9]),
                              position_error_m=float(np.linalg.norm(target[:3]-current[:3])),
                              ik_residual=[res_m, res_deg],
                              flags=[flags['clamped_workspace'], flags['limited_translation'],
                                     flags['limited_rotation']])
                steps += 1
            if a.mode == 'step':
                try:
                    input('Enter = next target, Ctrl-C = stop ')
                except EOFError:
                    break
    except KeyboardInterrupt:
        print('\nstopped by operator')
    finally:
        if a.mode != 'dry-run':
            arm.stop()
        cam.close()
        if rec:
            rec.close()
        print('%d steps executed' % steps)


# --------------------------------------------------------------------------- #
# selftest -- no robot, no camera, no policy server
# --------------------------------------------------------------------------- #
def selftest():
    ok = True

    def check(name, cond, detail=''):
        nonlocal ok
        ok = ok and bool(cond)
        print('%-42s %s %s' % (name, 'PASS' if cond else 'FAIL', detail))

    q = np.array([0.0, -0.3, 0.0, -2.0, 0.0, 1.8, 0.7])
    T = panda_fk(q)
    check('FK returns a proper rotation', np.abs(T[:3, :3].T @ T[:3, :3] - np.eye(3)).max() < 1e-9)
    check('FK det = +1', abs(np.linalg.det(T[:3, :3]) - 1) < 1e-9)

    # Jacobian against finite differences
    J, Jn, h = panda_jacobian(q), np.zeros((6, 7)), 1e-6
    T0 = panda_fk(q)
    for j in range(7):
        qp = q.copy()
        qp[j] += h
        T1 = panda_fk(qp)
        Jn[:3, j] = (T1[:3, 3] - T0[:3, 3]) / h
        dR = T1[:3, :3] @ T0[:3, :3].T
        Jn[3:, j] = np.array([dR[2, 1] - dR[1, 2], dR[0, 2] - dR[2, 0], dR[1, 0] - dR[0, 1]]) / (2 * h)
    check('Jacobian vs finite differences', np.abs(J - Jn).max() < 1e-5, '%.2e' % np.abs(J - Jn).max())

    # IK round trip
    p10 = pose10_from_matrix(T, 0.9)
    q2, rm, rd = solve_ik(p10, q + 0.05)
    check('IK round trip', rm < 2e-3 and rd < 0.6, '%.2f mm / %.3f deg' % (rm * 1000, rd))

    # pose10 round trip
    pos, R, quat, ap = unpack_pose10(p10)
    check('pose10 round trip', np.allclose(R, T[:3, :3], atol=1e-6) and abs(ap - 0.9) < 1e-6)

    # state_pose10
    s = state_pose10({'joint_positions': np.r_[q, 0.75]})
    check('state_pose10 shape/aperture', s.shape == (10,) and abs(s[9] - 0.75) < 1e-6)

    # pad_to_5x4
    p = pad_to_5x4(np.zeros((480, 640, 3), np.uint8))
    check('pad_to_5x4 640x480 -> 640x512', p.shape == (512, 640, 3), str(p.shape))
    check('pad_to_5x4 aspect is 5:4', abs(p.shape[1] / p.shape[0] - 1.25) < 1e-9)

    # PNG encode
    png = encode_png(np.zeros((8, 10, 3), np.uint8))
    check('PNG encode (magic bytes)', png[:8] == b'\x89PNG\r\n\x1a\n')

    # BGR->RGB
    f = np.zeros((2, 2, 3), np.uint8)
    f[..., 0] = 7
    check('bgr_to_rgb swaps channels', bgr_to_rgb(f)[0, 0, 2] == 7)

    # safety filter
    sf = SafetyFilter([0.2, -0.5, 0.0], [0.85, 0.5, 0.85])
    cur = pose10_from_matrix(T, 1.0)
    far = np.array(cur, np.float64)
    far[0] += 0.5
    check('check_chunk rejects a 0.5 m jump', not sf.check_chunk(far[None], cur)[0])
    check('check_chunk accepts the current pose', sf.check_chunk(np.asarray(cur)[None], cur)[0])
    tgt = np.array(cur, np.float64)
    tgt[0] += 0.10
    out, flags = sf.step(tgt, cur)
    check('step limits translation to 15 mm', flags['limited_translation']
          and abs(np.linalg.norm(out[:3] - cur[:3]) - 0.015) < 1e-6)
    hi = np.array(cur, np.float64)
    hi[2] = 9.0
    out2, flags2 = sf.step(hi, cur)
    check('step clamps to the workspace box', flags2['clamped_workspace'] and out2[2] <= 0.85 + 1e-6)
    out3, _ = sf.step(np.r_[cur[:9], 1.004], cur)
    check('step clips aperture into [0,1]', out3[9] <= 1.0)

    # gripper hysteresis
    g = GripperHysteresis(0.6, 0.4)
    check('gripper: no edge inside the dead band', g.update(0.5) is None)
    check('gripper: closes below 0.4', g.update(0.2) == 'close')
    check('gripper: no repeat while closed', g.update(0.3) is None)
    check('gripper: opens above 0.6', g.update(0.9) == 'open')
    check('gripper flag polarity (open = -1)', g.flag == GRIPPER_OPEN)

    check('all five prompts present', len(PROMPTS) == 5 and all(PROMPTS.values()))
    check('RECORD_ROOT is set', isinstance(RECORD_ROOT, str) and bool(RECORD_ROOT))
    hp = panda_fk(HOME_JOINTS)[:3, 3]
    check('HOME_JOINTS is 7 finite angles', HOME_JOINTS.shape == (7,) and np.isfinite(HOME_JOINTS).all())
    check('home is inside the joint limits',
          bool((HOME_JOINTS >= JOINT_LIMITS[:, 0]).all() and (HOME_JOINTS <= JOINT_LIMITS[:, 1]).all()))
    check('home flange is inside the workspace box',
          bool((hp >= np.array(DEMO_WORKSPACE[:3])).all() and (hp <= np.array(DEMO_WORKSPACE[3:])).all()),
          str(np.round(hp, 3)))
    check('reset aperture 1.0 opens the hand', (-1 if 1.0 >= 0.8 else 1) == -1)

    # --- gripper encoding, against gello/robots/panda.py's OWN rule ---
    def panda_binarise(v):
        """Verbatim gello/robots/panda.py:command_joint_state."""
        return -1 if v >= 0.8 else 1        # -1 = open (Move), +1 = close (Grasp)

    check('aperture 1.0 (open)  -> panda.py OPEN',  panda_binarise(1.0) == -1)
    check('aperture 0.0 (close) -> panda.py CLOSE', panda_binarise(0.0) == +1)
    check('the +-1 flag would INVERT on this API',
          panda_binarise(GRIPPER_OPEN) == +1 and panda_binarise(GRIPPER_CLOSE) == -1)
    # round trip through what we actually send
    for ap, want in ((1.0, -1), (0.95, -1), (0.0, +1), (0.2, +1)):
        got = panda_binarise(float(np.clip(ap, 0.0, 1.0)))
        check('send_joints aperture %.2f -> %s' % (ap, 'OPEN' if want == -1 else 'CLOSE'), got == want)

    # --- the delta path, against the pi0.5 eval script's own conversion ---
    Rf = panda_fk(q)[:3, :3]
    pf = panda_fk(q)[:3, 3]
    tp, tR = tcp_from_flange(pf, Rf)
    check('flange->TCP offset is 0.1034 m', abs(np.linalg.norm(tp - pf) - 0.1034) < 1e-9,
          '%.4f m' % np.linalg.norm(tp - pf))
    check('flange->TCP rotation is 45 deg', abs(geodesic_deg(Rf, tR) - 45.0) < 1e-6,
          '%.3f deg' % geodesic_deg(Rf, tR))
    d0 = abs_to_local_delta(tp, tR, tp, tR, GRIPPER_OPEN)
    check('zero delta when target == current',
          np.allclose(d0[:3], 0, atol=1e-12) and np.allclose(d0[3:9], matrix_to_rot6d(np.eye(3)), atol=1e-12))
    check('delta is 10 numbers, gripper last', d0.shape == (10,) and d0[9] == GRIPPER_OPEN)
    # a known local translation must come back in the CURRENT frame, not the base frame
    off = np.array([0.03, -0.01, 0.02])
    d1 = abs_to_local_delta(tp + tR @ off, tR, tp, tR, GRIPPER_CLOSE)
    check('delta_pos is expressed in the current TCP frame', np.allclose(d1[:3], off, atol=1e-12))
    # round trip: applying the delta to the current pose reproduces the target
    tgt_R = tR @ axis_angle_to_matrix(np.array([0.05, -0.02, 0.01]))
    d2 = abs_to_local_delta(tp + tR @ off, tgt_R, tp, tR, GRIPPER_OPEN)

    def dp_gs_rot_6d_to_rot_mat(r6):
        """Verbatim dp_gs.dataset.utils.rot_6d_to_rot_mat: rows 0,1 then cross(row0,row1)."""
        v = np.asarray(r6, np.float64).reshape(2, 3)
        M = np.zeros((3, 3))
        M[:2, :] = v
        M[2, :] = np.cross(v[0], v[1])
        return M

    check('delta rot6d uses the ROW convention (dp_gs), not JAM columns',
          np.allclose(d2[3:9], (tR.T @ tgt_R)[:2, :].reshape(6), atol=1e-12))
    check('delta rotation round trip through dp_gs own inverse',
          np.allclose(tR @ dp_gs_rot_6d_to_rot_mat(d2[3:9]), tgt_R, atol=1e-12))
    check('row and column rot6d genuinely differ (trap is real)',
          not np.allclose(matrix_to_rot6d_rows(tgt_R), matrix_to_rot6d(tgt_R)))
    check('quat_xyzw_to_matrix round trip',
          np.allclose(quat_xyzw_to_matrix(matrix_to_quat_xyzw(tR)), tR, atol=1e-9))
    check('tcp_from_obs reads eef_pos/eef_quat',
          np.allclose(tcp_from_obs({'eef_pos': tp, 'eef_quat': matrix_to_quat_xyzw(tR)})[0], tp))
    import re as _re, tempfile as _tf, os as _os
    with _tf.TemporaryDirectory() as td:
        ds = [Recorder(td, 'task_1', '123', attrs={}).dir for _ in range(5)]
        names = [_os.path.basename(x) for x in ds]
        check('session name matches the demos (session_<8 hex>)',
              all(_re.fullmatch(r'session_[0-9a-f]{8}', n) for n in names), names[0])
        check('session names are unique', len(set(names)) == 5)
        check('nested under <root>/task_1/',
              all(_os.path.basename(_os.path.dirname(x)) == 'task_1' for x in ds))
    lo, hi = np.array(DEMO_WORKSPACE[:3]), np.array(DEMO_WORKSPACE[3:])
    check('default workspace is a valid box', (hi > lo).all())
    demo_lo, demo_hi = np.array([0.457, -0.220, 0.211]), np.array([0.746, 0.356, 0.548])
    check('default box contains the demo range', (lo < demo_lo).all() and (hi > demo_hi).all())
    sf2 = SafetyFilter(DEMO_WORKSPACE[:3], DEMO_WORKSPACE[3:])
    check('default box accepts a mid-demo pose', not sf2.step(
        np.r_[0.6, 0.05, 0.35, matrix_to_rot6d(np.eye(3)), 1.0],
        np.r_[0.6, 0.05, 0.35, matrix_to_rot6d(np.eye(3)), 1.0])[1]['clamped_workspace'])
    print('\n%s' % ('SELFTEST PASSED' if ok else 'SELFTEST FAILED'))
    return 0 if ok else 1


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--selftest', action='store_true', help='check the maths with no hardware and exit')
    p.add_argument('--task', type=int, choices=[1, 2, 3, 4, 5], help='selects the baked-in prompt')
    p.add_argument('--prompt', help='override; must match the adapter VERBATIM')
    p.add_argument('--url', default='http://127.0.0.1:8000/act')
    p.add_argument('--mode', choices=['dry-run', 'step', 'run'], default='dry-run')
    p.add_argument('--execute-steps', type=int, default=8, help='targets executed per chunk (1..32)')
    p.add_argument('--control', choices=['joint', 'delta'], default='joint',
                   help="joint (default): IK to absolute joints -> command_joint_state, the same "
                        "call GELLO teleoperation uses and the exact inverse of how the training "
                        "targets were built. delta: TCP-frame command_delta_pose, the extension "
                        "the pi0.5 evals added for their EE-space policy")
    p.add_argument('--workspace', type=float, nargs=6, metavar=('X0', 'Y0', 'Z0', 'X1', 'Y1', 'Z1'),
                   help='safe box for the FLANGE, metres. Default is the demonstrations own '
                        'bounding box + 5 cm: %s' % DEMO_WORKSPACE)
    p.add_argument('--max-step-m', type=float, default=0.015)
    p.add_argument('--max-step-deg', type=float, default=4.0)
    p.add_argument('--gripper-close-below', type=float, default=0.4)
    p.add_argument('--gripper-open-above', type=float, default=0.6)
    p.add_argument('--max-steps', type=int, default=300,
                   help='budget in EXECUTED targets (15 Hz), not seconds')
    p.add_argument('--seed', type=int, default=0, help='fixed diffusion seed across replans for repeatable diagnosis')
    p.add_argument('--timeout', type=float, default=120.0, help='HTTP timeout; the first call is slow')
    p.add_argument('--camera-serial', default=HEAD_SERIAL)
    p.add_argument('--robot-host', default='127.0.0.1')
    p.add_argument('--robot-port', type=int, default=6001)
    p.add_argument('--record', default=RECORD_ROOT,
                   help='root for task_<N>/session_<8 hex>/ (h5 + mp4). Default: %s' % RECORD_ROOT)
    p.add_argument('--no-record', action='store_true', help='do not write the h5/mp4 rollout')
    p.add_argument('--no-reset', action='store_true',
                   help='skip the move to the demonstrations start pose (step/run only; '
                        'dry-run never resets)')
    a = p.parse_args()
    if a.selftest:
        sys.exit(selftest())
    if a.task is None and not a.prompt:
        sys.exit('give --task 1..5 (or --prompt). Run --selftest first.')
    if not 1 <= a.execute_steps <= 32:
        sys.exit('--execute-steps must be within 1..32')
    run(a)


if __name__ == '__main__':
    main()