"""deploy_with_action_chunking.py — deploy.py + ACT-style temporal ensembling. Only difference vs. ``deploy.py``: when re-querying mid-rollout, this script log-pools the per-bin **logits** of the last ``--ensemble_chunks`` chunks before decoding. Adjacent chunks' overlapping timesteps each contribute exp(-m · chunks_ago)-weighted contributions; the decode argmax then picks the bin that all recent chunks agree on. With ``--ensemble_chunks 1`` (default) the script is bit-identical to ``deploy.py``. We pool **logits** (not decoded xyz/grip/quat) because PARA's heads are categorical — log-pooling preserves multimodality and unifies the xyz / grip / rot heads under one operation (no SLERP gymnastics for quaternions). See: - Zhao et al. 2023 — ACT/ALOHA temporal-ensembling (continuous version) - Hinton 2002 — Product-of-Experts / log-pooling distributions Original deploy.py docstring follows. deploy.py — live inference + rerun + raiden control on russet. Direct-raiden inference (russet has a 4090 → no chiral roundtrip needed). Same rerun viz scaffold as ``calibrate.py`` / ``record.py``; same ``Robot.move_to`` wrapper for the joint-sequence execution that ``control_speed_test.py`` exercises. Each iteration: 1. Grab scene RGB + read current 7-DoF right-arm joints from the follower. 2. FK on joints → project to scene image → ``start_pix`` (model input). 3. Forward model → ``volume_logits`` + grip + rot. 4. Decode (pixel + height_bin + rot_bin + grip_bin) → world XYZ + quaternion + gripper value per future timestep. 5. mink-IK each EE pose (seeded chain) → ``(T, 7)`` joint commands. 6. Log rerun panels (live RGB, mujoco-overlay at current joints, pred keypoints over live, heatmap_grid with + without RGB, status). 7. Prompt y / q / h / r / c. ``y`` runs the chunk via ``Robot.move_to``. Run on russet:: ssh robot-lab source ~/cameron/raiden_fork.venv/bin/activate DINO_WEIGHTS_PATH=$HOME/cameron/puget/dinov3/weights/\\ dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth \\ python ~/lab/para/robot/yam/deploy.py \\ --ckpt $HOME/cameron/puget/checkpoints/place_mug_saucer_v0/latest.pth """ from __future__ import annotations import argparse import json import math import os import sys import time from datetime import datetime from pathlib import Path os.environ.setdefault("MUJOCO_GL", "egl") # Auto-route DINOv3 weights through ~/cameron/puget mount when on russet. _puget_w = Path.home() / "cameron/puget/dinov3/weights/" \ "dinov3_vits16plus_pretrain_lvd1689m-4057cbaa.pth" if _puget_w.exists(): os.environ.setdefault("DINO_WEIGHTS_PATH", str(_puget_w)) import cv2 import numpy as np import rerun as rr import torch sys.path.insert(0, str(Path(__file__).resolve().parent)) from lib.calibration import pool_exo_aruco, load_T_lfr, to_world from lib.deploy_panels import (_log_scene_overlay, log_baseline_panels, log_volume_panels, log_volume_panels_multiview, log_volume_panels_v3) from lib.ik import ee_chain_to_joints from lib.inference import ( load_checkpoint, decode as decode_outputs, decode_baseline, decode_baseline_motion_tracks, decode_volumetric, decode_volumetric_multiview, decode_v3) from lib.rerun_helpers import init_web from lib.robot import Robot from lib.video_recorder import VideoRecorder from lib.dataset import IMAGENET_MEAN, IMAGENET_STD, _fk_ee_in_rbase CAM_CFG_PATH = Path.home() / ".config" / "raiden" / "camera.json" WRIST_CAL_PATH = Path.home() / ".config" / "raiden" / "wrist_calibration_result.json" def _open_scene_zed(name: str = "scene_camera"): """Open the ZED registered under ``name`` in ``~/.config/raiden/camera.json``. Pass ``"scene_camera_ood"`` to use the third ZED 2i for OOD-viewpoint eval sessions (added 2026-06-11).""" from raiden.calibration.exo_calibrate import _open_zed_native serial = int(json.load(open(CAM_CFG_PATH))[name]["serial"]) cam, K, _dist, W, H = _open_zed_native(serial, "HD1080") return cam, np.asarray(K, dtype=np.float64), int(W), int(H) def _open_wrist_zed(arm: str): """Open the per-arm wrist ZED (same pattern as scene). Returns (cam, K_native, W, H). Native res = HD1080 (matches the training-time intrinsics persisted in wrist_calibration_result.json).""" from raiden.calibration.exo_calibrate import _open_zed_native cam_cfg = json.load(open(CAM_CFG_PATH))[f"{arm}_wrist_camera"] cam, K, _dist, W, H = _open_zed_native(int(cam_cfg["serial"]), "HD1080") return cam, np.asarray(K, dtype=np.float64), int(W), int(H) def _load_wrist_hand_eye(arm: str) -> np.ndarray: """4×4 T_ee_wrist_cam from raiden's saved wrist calibration.""" d = json.load(open(WRIST_CAL_PATH)) he = d["cameras"][f"{arm}_wrist_camera"] he = he.get("hand_eye_calibration") or he.get("extrinsics") R = np.asarray(he["rotation_matrix"], dtype=np.float64) t = np.asarray(he["translation_vector"], dtype=np.float64).reshape(3) T = np.eye(4); T[:3, :3] = R; T[:3, 3] = t return T def _grab_bgr(cam): import pyzed.sl as sl if cam.grab(sl.RuntimeParameters()) != sl.ERROR_CODE.SUCCESS: return None img = sl.Mat() cam.retrieve_image(img, sl.VIEW.LEFT) return cv2.cvtColor(img.get_data(), cv2.COLOR_BGRA2BGR) def _scale_K(K, src_wh, dst_wh): K2 = K.copy() K2[0, 0] *= dst_wh[0] / src_wh[0]; K2[0, 2] *= dst_wh[0] / src_wh[0] K2[1, 1] *= dst_wh[1] / src_wh[1]; K2[1, 2] *= dst_wh[1] / src_wh[1] return K2 def _imagenet_norm(bgr, img_size): rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB) rgb = cv2.resize(rgb, (img_size, img_size), interpolation=cv2.INTER_AREA) x = (rgb.astype(np.float32) / 255.0 - IMAGENET_MEAN) / IMAGENET_STD return torch.from_numpy(x.transpose(2, 0, 1).copy()) def _project(pt_world, K, T_w2c): p = T_w2c[:3, :3] @ pt_world + T_w2c[:3, 3] uv = K @ p return uv[:2] / max(uv[2], 1e-6) def main(): p = argparse.ArgumentParser() p.add_argument("--ckpt", required=True) p.add_argument("--arm", default="right", choices=["right", "left"]) p.add_argument("--rerun_port", type=int, default=9092) p.add_argument("--n_detections", type=int, default=10) p.add_argument("--max_joint_vel_rad_s", type=float, default=0.3) p.add_argument("--first_move_vel_rad_s", type=float, default=0.2) p.add_argument("--render_scale", type=float, default=0.5) p.add_argument("--device", default="auto", help="'auto' picks cuda if available, else cpu. " "russet's raiden venv currently runs on cpu (~1-2s/forward).") p.add_argument("--save_video", action="store_true", help="Save one MP4 per rollout (h-boundary) under " "--eval_videos_root/_/.") p.add_argument("--eval_name", default="eval", help="Dirname prefix for the eval session " "(e.g. 'ours_in_dist_pickplace'). A timestamp is " "appended automatically.") p.add_argument("--eval_videos_root", default=str(Path.home() / "cameron/eval_videos"), help="Where the eval session directory goes (default = " "russet-local ~/cameron/eval_videos/, since the puget " "sshfs mount at ~/cameron/puget/ is read-only). " "Rsync to puget after the session if archival is needed.") p.add_argument("--video_fps", type=int, default=5, help="Playback fps for the recorded MP4. 5 ≈ a stop-motion " "speed; raise for slower playback.") p.add_argument("--t_lfr_path", default=None, help="Path to a T_left_from_right .npy file. Required ONLY " "for legacy ckpts trained on world=left_arm_base data " "(e.g. cal/T_left_from_right.npy). New ckpts trained on " "world=right_arm_base data (yukon, post-2026-06-04) " "leave this unset → T_lfr defaults to identity.") p.add_argument("--n_substeps", type=int, default=None, help="Execute only the first N substeps of each predicted " "chunk before replanning (receding-horizon style). " "Default = None → execute the full T-step chunk.") p.add_argument("--ensemble_chunks", type=int, default=1, help="ACT-style temporal ensembling depth (max number of " "past chunks whose predictions are log-pooled with the " "current chunk's predictions). 1 = no smoothing — " "identical to deploy.py. 4 ≈ T/n_substeps for the " "default T=32, n_substeps=8.") p.add_argument("--use_alt_cam", action="store_true", help="Open ``scene_camera_ood`` (third ZED 2i, OOD-viewpoint " "testing cam, added 2026-06-11) instead of the canonical " "``scene_camera``. Re-locks via aruco per session — no " "persisted pose so each placement is treated fresh.") p.add_argument("--border_mask_px", type=int, default=5, help="Zero a frame of this width around each view's " "pred-grid edge before argmax. Prevents OOD object " "placements from latching onto the image border, " "where features have low training supervision. " "Default 5, set 0 to disable.") p.add_argument("--ensemble_decay", type=float, default=0.05, help="Exponential decay constant m for chunk weights: " "w_i = exp(-m * chunks_ago). ACT uses 0.01 (very gentle " "— old chunks contribute meaningfully). Default 0.05 " "since we have fewer overlapping chunks than ACT's " "K=100 setup.") p.add_argument("--z_offset_m", type=float, default=0.0, help="Hacky world-Z shift applied to the decoded EE xyz " "BEFORE IK. Negative values pull predictions toward " "the table; e.g. -0.0127 ≈ half an inch down. " "Use to A/B the 'always lands above the table' bias " "rooted in (a) sparse low-z training data and (b) " "limited sin/cos PE resolution at height_enc_dim=16.") args = p.parse_args() device = args.device if device == "auto": device = "cuda" if torch.cuda.is_available() else "cpu" ck = load_checkpoint(Path(args.ckpt), device=device) args.device = device # save for later passes model = ck["model"]; cfg = ck["cfg"]; kind = ck["kind"] is_baseline_mt = (kind == "baseline_cls_uvz_motion_tracks") is_baseline_diff = (kind == "baseline_cls_xyz_diffusion") is_baseline = (kind in ("baseline_cls_xyz", "baseline_cls_xyz_multiview", "baseline_cls_xyz_diffusion", "baseline_cls_uvz_motion_tracks")) is_baseline_mv = (kind == "baseline_cls_xyz_multiview") is_v3 = (kind == "dino_volume_scene_v3") is_v4 = (kind in ("dino_volume_scene_v4", "dino_volume_scene_v4_da3", "dino_volume_scene_v4_vlm", "dino_volume_scene_v4_dynaflip", "dino_volume_scene_v4_paligemma", "dino_volume_scene_v4_pi0")) is_multiview = (kind == "dino_volume_scene_volumetric_multiview" or is_v3 or is_v4) # ``views`` comes from either a multi-view volumetric model or the # multi-view baseline — both have a .views attribute. views = (list(model.views) if (is_multiview or is_baseline_mv) else ["scene"]) has_scene_view = "scene" in views has_wrist_view = "wrist" in views # Past-trajectory proprio conditioning (only when the multi-view model # was trained with --mv_past_n > 0). We keep the most-recent # ``mv_past_n`` (grip, world-Z) pairs and pass them on every forward. mv_past_n = int(getattr(model, "past_n", 0)) if is_multiview else 0 past_grip_deque: list[float] = [] past_z_deque: list[float] = [] print(f"loaded ckpt step={ck['step']} kind={kind} from {args.ckpt} (device={device})") if is_multiview: print(f" multi-view: views={views} primary={views[0]}") scene_cam_name = "scene_camera_ood" if args.use_alt_cam else "scene_camera" if args.use_alt_cam: print(f" ⚑ --use_alt_cam: opening '{scene_cam_name}' (OOD viewpoint cam)") cam, K_native, W, H = _open_scene_zed(name=scene_cam_name) init_web("yam_deploy", port=args.rerun_port) print(f"scene cam {W}x{H} K diag=({K_native[0,0]:.1f}, {K_native[1,1]:.1f})") # Wrist cam + hand-eye — only when the model consumes the wrist view. cam_wrist = None K_wrist_native = None Ww = Hw = 0 T_ee_wrist_cam = None if has_wrist_view: cam_wrist, K_wrist_native, Ww, Hw = _open_wrist_zed(args.arm) T_ee_wrist_cam = _load_wrist_hand_eye(args.arm) print(f"wrist cam {Ww}x{Hw} K diag=({K_wrist_native[0,0]:.1f}, " f"{K_wrist_native[1,1]:.1f})") print(f" hand-eye t = {T_ee_wrist_cam[:3, 3].round(4).tolist()}") # T_lfr: identity for the new world=right_arm_base convention (yukon-era # ckpts); load from --t_lfr_path for legacy world=left_arm_base ckpts. if args.t_lfr_path: T_lfr = load_T_lfr(Path(args.t_lfr_path)) print(f" T_lfr loaded from {args.t_lfr_path} (legacy world=lbase)") else: T_lfr = np.eye(4, dtype=np.float64) print(" T_lfr = identity (world=right_arm_base; new convention)") robot = Robot(use_left=(args.arm == "left"), use_right=(args.arm == "right")) follower = robot.follower_r if args.arm == "right" else robot.follower_l home_q7 = follower.get_joint_pos().copy() # treat current as home print(f"robot ready; home (current) q7 = {home_q7}") # Lock cam pose. print(f"locking cam pose: pooling {args.n_detections} aruco frames…") T_cam_in_rbase = pool_exo_aruco( lambda: (_grab_bgr(cam), K_native), min_detections=args.n_detections, ) T_cam_in_world = to_world(T_cam_in_rbase, T_lfr) T_world_to_cam = np.linalg.inv(T_cam_in_world) print(f"✓ locked. cam t = {T_cam_in_rbase[:3, 3]}") img_size = cfg["img_size"] K_in = _scale_K(K_native, (W, H), (img_size, img_size)) Wr, Hr = max(1, int(W * args.render_scale)), max(1, int(H * args.render_scale)) K_render = _scale_K(K_native, (W, H), (Wr, Hr)) K_in_wrist = (_scale_K(K_wrist_native, (Ww, Hw), (img_size, img_size)) if cam_wrist is not None else None) # Startup calibration overlay: render mujoco arm @ current joints on the # live scene, log scene/{rgb,composite,mask_overlay} so the user can # eyeball cam-vs-robot alignment before running any inference. _bgr0 = _grab_bgr(cam) _joints0 = np.asarray(follower.get_joint_pos(), dtype=np.float64) if _bgr0 is not None: rr.set_time("step", sequence=0) _log_scene_overlay(_bgr0, _joints0, T_cam_in_rbase, K_render, Wr, Hr, W, H) print(" ✓ logged calibration overlay → scene/composite (eyeball it before pressing y/r)") # Eval-session dir is always created (post-2026-06-22): per-chunk dumps # (input RGBs + camera poses + intrinsics + joints7 + decoded XYZ/rot/ # grip + small logits + high-res patch features) are produced for every # rollout so eval analyses can re-render PCA / pred-trajectory viz # off-line without re-running inference. --save_video is now a strict # add-on for the MP4 recording only. ts = datetime.now().strftime("%Y-%m-%d_%H%M%S") eval_dir = Path(args.eval_videos_root) / f"{args.eval_name}_{ts}" eval_dir.mkdir(parents=True, exist_ok=True) print(f" eval session dir → {eval_dir}/ (per-chunk dumps always on)") recorder = None if args.save_video: recorder = VideoRecorder(eval_dir, W=W, H=H, fps=args.video_fps) print(f" recording MP4 rollouts → {eval_dir}/") # Start frame of vid_001 (current scene right after cam lock). bgr0 = _grab_bgr(cam) if bgr0 is not None: recorder.append(bgr0) frame_idx = 0 rollout_idx = 1 # Auto-yes loop: after the first y press, subsequent chunks execute # automatically without re-prompting. Ctrl-C during a chunk pauses # the robot and drops back to manual prompt mode. auto_yes = False print(f"\n ── rollout {rollout_idx} ──") if is_baseline_mt and args.ensemble_chunks > 1: print(" WARNING: motion-tracks baseline regresses (u,v,z) — chunk " "ensembling on raw_out would need a UVZ-aware pooler; capping " "--ensemble_chunks at 1.") args.ensemble_chunks = 1 if is_v3 and args.ensemble_chunks > 1: print(f" v3: action-chunk ensembling not yet implemented for the " f"factorized model — capping --ensemble_chunks at 1.") args.ensemble_chunks = 1 if args.ensemble_chunks > 1: print(f" action-chunk ensembling: depth={args.ensemble_chunks}, " f"decay m={args.ensemble_decay} (log-pooled per-bin logits)") print("\n[y]=deploy / [r]=re-query / [h]=home (stay) / " "[c]=recalibrate cam / [q]=quit+home") # ACT-style temporal ensembling state. Each entry is # (abs_start_step, raw_out_dict). Length capped at args.ensemble_chunks. # The number of substeps actually executed between chunks (used as the # offset stride) is whatever len(joint_cmds) actually runs — we record # it at execution time. chunk_buffer: list = [] last_stride = 0 # substeps executed since the chunk now at top-of-buffer abs_step_counter = 0 def _log_pool_raw_out_baseline(stride_per_chunk: int) -> dict: """Baseline variant: ``xyz`` is a continuous regression — averaged in Euclidean space (per Cameron's spec, ACT-style). ``grip_logits`` and ``rot_logits`` are still log-pooled (they're classifiers). Same offset arithmetic as the volume version: newest chunk's timestep i ↔ chunk_k's timestep (i + k*stride).""" newest_raw = chunk_buffer[-1][1] T_window = int(newest_raw["xyz_pred"].shape[1]) keys_logpool = tuple(k for k in ("grip_logits", "rot_logits") if k in newest_raw) accum_xyz = torch.zeros_like(newest_raw["xyz_pred"]) accum_log = {k: torch.zeros_like(newest_raw[k]) for k in keys_logpool} weight_per_t = torch.zeros(T_window, device=newest_raw["xyz_pred"].device, dtype=newest_raw["xyz_pred"].dtype) for chunks_ago, (_t0, raw, _K, _Tcw) in enumerate(reversed(chunk_buffer)): shift = chunks_ago * stride_per_chunk if shift >= T_window: break w = math.exp(-args.ensemble_decay * chunks_ago) accum_xyz[:, : T_window - shift] = ( accum_xyz[:, : T_window - shift] + w * raw["xyz_pred"][:, shift:T_window]) for k in keys_logpool: accum_log[k][:, : T_window - shift] = ( accum_log[k][:, : T_window - shift] + w * raw[k][:, shift:T_window]) weight_per_t[: T_window - shift] = weight_per_t[: T_window - shift] + w w_safe = weight_per_t.clamp(min=1e-12) out = {"xyz_pred": accum_xyz / w_safe.view(1, T_window, 1)} for k in keys_logpool: out[k] = accum_log[k] / w_safe.view(1, T_window, *([1] * (accum_log[k].dim() - 2))) for k, v in newest_raw.items(): if k not in out: out[k] = v return out def _log_pool_raw_out(stride_per_chunk: int) -> dict: """Build an ensembled raw_out by log-pooling per-bin logits across chunks in ``chunk_buffer``. Newer chunks weigh more (exp decay). For chunk at index k (counted from newest = 0), its prediction for ABSOLUTE timestep T_abs = chunk[k].abs_start + offset_k. Given the newest chunk's offset i, we want abs_step = newest.abs_start + i, so older chunk k's offset is i + k * stride_per_chunk. Only chunks whose offset stays inside the model's prediction window (< T_window) contribute. Beyond that horizon a chunk's prediction was never produced. """ # All chunks share the same shape (T_window, ...) for each logit head. newest_raw = chunk_buffer[-1][1] T_window = int(newest_raw["volume_logits"].shape[1]) # Initialise zero accumulators on the same device/dtype. keys = ("volume_logits", "grip_logits", "rot_logits") accum = {k: torch.zeros_like(newest_raw[k]) for k in keys} weight_per_t = torch.zeros(T_window, device=newest_raw["volume_logits"].device, dtype=newest_raw["volume_logits"].dtype) # Iterate newest → oldest so chunks_ago grows. for chunks_ago, (_abs_start, raw, _K, _Tcw) in enumerate(reversed(chunk_buffer)): shift = chunks_ago * stride_per_chunk if shift >= T_window: break w = math.exp(-args.ensemble_decay * chunks_ago) for k in keys: # Newest chunk's offset i ↔ this chunk's offset i+shift. # So we add raw[k][:, shift:T_window] into accum[k][:, : T_window-shift] # weighted by w. accum[k][:, : T_window - shift] = ( accum[k][:, : T_window - shift] + w * raw[k][:, shift:T_window]) weight_per_t[: T_window - shift] = weight_per_t[: T_window - shift] + w # Normalise (avoid div-by-zero by clamping). w_safe = weight_per_t.clamp(min=1e-12) out = {k: accum[k] / w_safe.view(1, T_window, *([1] * (accum[k].dim() - 2))) for k in keys} # Preserve any other tensors the panel viz expects (e.g. F_refined, # patch_features) — viz just uses the current chunk's features. for k, v in newest_raw.items(): if k not in out: out[k] = v return out def _reverse_project_to_old_indices(new_pos_flat, old_K_list, old_T_cw_list, n_views, Z, P, img_size_px, z_lo_, z_hi_): """Closed-form O(V) mapping from newest chunk's voxel positions to the OLDER chunk's flat voxel index, slab by slab. New voxel k*Z.. (k+1)*Z stays in slab k of the old grid; we project its world XYZ through OLDER view k's camera and snap to (z_bin, yi, xi).""" V = new_pos_flat.shape[0] cell = img_size_px / P out_idx = np.empty(V, dtype=np.int64) slab_of = np.arange(V) // (Z * P * P) slab_size = Z * P * P for k in range(n_views): sel = (slab_of == k) pts = new_pos_flat[sel] # (Vk, 3) T_w2c = np.linalg.inv(old_T_cw_list[k]).astype(np.float64) K = np.asarray(old_K_list[k], dtype=np.float64) cam = pts @ T_w2c[:3, :3].T + T_w2c[:3, 3] # (Vk, 3) z_cam = np.clip(cam[:, 2], 1e-3, None) u = (K[0, 0] * cam[:, 0] / z_cam) + K[0, 2] v = (K[1, 1] * cam[:, 1] / z_cam) + K[1, 2] xi = np.clip(np.round(u / cell - 0.5), 0, P - 1).astype(np.int64) yi = np.clip(np.round(v / cell - 0.5), 0, P - 1).astype(np.int64) pz = pts[:, 2] zb = np.clip(np.round((pz - z_lo_) / (z_hi_ - z_lo_) * Z - 0.5), 0, Z - 1).astype(np.int64) out_idx[sel] = k * slab_size + zb * (P * P) + yi * P + xi return out_idx def _log_pool_volumetric_multiview(stride_per_chunk: int) -> dict: """Volume-space log-pool via closed-form reverse projection. Each chunk's voxel positions live in a different world frame (wrist cam moves between chunks); we map every newest-chunk voxel into each older chunk's grid by projecting its world XYZ through the older chunk's per-view camera + snapping (z_bin, yi, xi). O(V) per chunk vs O(V·log V) for the cKDTree path it replaces. Grip + rot are straight log-pool (bin defs don't move).""" newest_t0, newest_raw, _newK, _newT = chunk_buffer[-1] new_logits = newest_raw["volume_logits"][0] # (T, NZ, P, P) new_pos = newest_raw["voxel_positions"][0] # (NZ, P, P, 3) T_w, NZ, P, _ = new_logits.shape V = NZ * P * P n_v = len(views) Z = NZ // n_v new_pos_flat = new_pos.reshape(V, 3).detach().cpu().numpy() z_lo_, z_hi_ = cfg["z_range"] accum_vol = torch.zeros_like(new_logits) accum_grip = torch.zeros_like(newest_raw["grip_logits"]) accum_rot = torch.zeros_like(newest_raw["rot_logits"]) weight_per_t = torch.zeros(T_w, device=new_logits.device, dtype=new_logits.dtype) for chunks_ago, (_t0, raw, old_K, old_Tcw) in enumerate(reversed(chunk_buffer)): shift = chunks_ago * stride_per_chunk if shift >= T_w: break w = math.exp(-args.ensemble_decay * chunks_ago) old_flat = raw["volume_logits"][0].reshape(T_w, V) if chunks_ago == 0: mapped_flat = old_flat else: idx = _reverse_project_to_old_indices( new_pos_flat, old_K, old_Tcw, n_v, Z, P, img_size, z_lo_, z_hi_) idx_t = torch.from_numpy(idx).long().to(old_flat.device) mapped_flat = old_flat[:, idx_t] mapped = mapped_flat.view(T_w, NZ, P, P) accum_vol[: T_w - shift] = ( accum_vol[: T_w - shift] + w * mapped[shift:T_w]) accum_grip[:, : T_w - shift] = ( accum_grip[:, : T_w - shift] + w * raw["grip_logits"][:, shift:T_w]) accum_rot[:, : T_w - shift] = ( accum_rot[:, : T_w - shift] + w * raw["rot_logits"][:, shift:T_w]) weight_per_t[: T_w - shift] = weight_per_t[: T_w - shift] + w w_safe = weight_per_t.clamp(min=1e-12) out = dict(newest_raw) out["volume_logits"] = (accum_vol / w_safe.view(T_w, 1, 1, 1)).unsqueeze(0) out["grip_logits"] = accum_grip / w_safe.view(1, T_w, 1) out["rot_logits"] = accum_rot / w_safe.view(1, T_w, 1) return out def go_home(): """Smooth-move back to the joints we read at startup.""" target14 = (np.concatenate([np.zeros(7), home_q7]) if args.arm == "right" else np.concatenate([home_q7, np.zeros(7)])) robot.move_to(target14, vel_rad_s=args.first_move_vel_rad_s) class _RolloutSaver: """Saves per-chunk artifacts under the same eval_dir as the MP4. Layout: ``/rollout_NNN/chunk_MMM/`` with: - frame_pre.jpg input image at the moment of decision - predictions.npz decoded outputs (xyz, quat, grip, bins …) + small grip/rot logits (no volume — too big) - patch_features.npy raw DINO 31×31 patch grid (fp16, ~1.5 MB) - camera_pose.npy T_cam_in_rbase (4×4) - substep_NN.jpg frame after each executed substep """ def __init__(self, eval_dir: Path): self.eval_dir = eval_dir self.rollout_idx = 1 self.chunk_idx = 0 def rotate(self): self.rollout_idx += 1 self.chunk_idx = 0 def chunk_dir(self) -> Path: d = (self.eval_dir / f"rollout_{self.rollout_idx:03d}" / f"chunk_{self.chunk_idx:03d}") d.mkdir(parents=True, exist_ok=True) return d def save_chunk(self, decoded: dict, bgr_pre: np.ndarray, T_cam_in_rbase: np.ndarray, bgr_pre_wrist: np.ndarray | None = None, T_wrist_in_world: np.ndarray | None = None, K_scene_img_size: np.ndarray | None = None, K_wrist_img_size: np.ndarray | None = None, joints7: np.ndarray | None = None): d = self.chunk_dir() cv2.imwrite(str(d / "frame_pre.jpg"), bgr_pre, [cv2.IMWRITE_JPEG_QUALITY, 90]) np.save(d / "camera_pose.npy", T_cam_in_rbase.astype(np.float64)) # Wrist view (added 2026-06-12). Wrist BGR + per-step pose so the # figure-maker / post-hoc analysis pipeline has both views. if bgr_pre_wrist is not None: cv2.imwrite(str(d / "frame_pre_wrist.jpg"), bgr_pre_wrist, [cv2.IMWRITE_JPEG_QUALITY, 90]) if T_wrist_in_world is not None: np.save(d / "camera_pose_wrist.npy", T_wrist_in_world.astype(np.float64)) # Intrinsics scaled to the model's img_size — same matrix every chunk # but cheap to repeat per chunk and keeps each chunk self-contained. if K_scene_img_size is not None: np.save(d / "K_scene_img_size.npy", K_scene_img_size.astype(np.float64)) if K_wrist_img_size is not None: np.save(d / "K_wrist_img_size.npy", K_wrist_img_size.astype(np.float64)) # Proprio input — the 7-DoF joint vector the model implicitly # conditioned on (via FK→start_pix or past-trajectory deque). if joints7 is not None: np.save(d / "joints7.npy", np.asarray(joints7, dtype=np.float64)) # Predictions — strip the huge ``raw_out`` (volume_logits is 1 GB). # Keep small logits + small raw heads (uvz_pred for motion-tracks, # xyz_pred for the regression baselines) so post-hoc viz can show # the model's pre-lift / pre-argmax prediction. raw = decoded.get("raw_out", {}) pred: dict = {} for k, v in decoded.items(): if k == "raw_out": continue if isinstance(v, np.ndarray): pred[k] = v for k in ("grip_logits", "rot_logits", "uvz_pred", "xyz_pred"): if k in raw: v = raw[k] if torch.is_tensor(v): v = v[0].detach().cpu().numpy() pred[k] = v np.savez_compressed(d / "predictions.npz", **pred) # Per-view refined feature maps (added 2026-06-12). For each view # save (F, P, P) fp16 — typically ~1 MB/view/chunk. Keeps both views # available for post-hoc PCA / heatmap visualisations. if "view_feat_maps" in raw: view_dict = {} for k, v in enumerate(raw["view_feat_maps"]): view_dict[f"view_{k}"] = v[0].detach().cpu().numpy().astype(np.float16) np.savez_compressed(d / "view_feat_maps.npz", **view_dict) # Lower-res DINO patches if the model exposes them. if "patch_features" in raw: pf = raw["patch_features"][0].detach().cpu().numpy() np.save(d / "patch_features.npy", pf.astype(np.float16)) def save_substep_frame(self, bgr: np.ndarray, substep_idx: int): cv2.imwrite(str(self.chunk_dir() / f"substep_{substep_idx:02d}.jpg"), bgr, [cv2.IMWRITE_JPEG_QUALITY, 90]) def end_chunk(self): self.chunk_idx += 1 rollout_saver = _RolloutSaver(eval_dir) class _FixedRawOutModel: """Stub that returns a pre-computed raw_out dict from ``__call__``. Lets us reuse ``decode_outputs`` from lib/inference.py with our log-pooled ensembled logits — no need to duplicate the workspace mask / border mask / argmax / unproject pipeline.""" def __init__(self, raw): self.raw = raw def __call__(self, *_args, **_kwargs): return self.raw try: while True: try: bgr = _grab_bgr(cam) if bgr is None: print(" cam grab failed; retrying"); continue joints7 = np.asarray(follower.get_joint_pos(), dtype=np.float64) rgb_t = _imagenet_norm(bgr, img_size).unsqueeze(0).to(args.device) if is_baseline: # Baseline directly regresses (T, 3) world XYZ; no start_pix needed. # Multi-view baseline takes a list of per-view RGBs; single-view # takes a single tensor. decode_baseline just forwards to model. if is_baseline_mv: bgr_wrist_b = (_grab_bgr(cam_wrist) if cam_wrist is not None else None) if has_wrist_view and bgr_wrist_b is None: print(" wrist cam grab failed; retrying"); continue rgb_for_model = [] for v in views: if v == "scene": rgb_for_model.append(rgb_t) elif v == "wrist": rgb_for_model.append(_imagenet_norm(bgr_wrist_b, img_size) .unsqueeze(0).to(args.device)) else: rgb_for_model = rgb_t if is_baseline_mt: # Motion-tracks regresses (u, v, world-z) in the scene-cam # image; lift to world XYZ via the live scene-cam intrinsics # + extrinsics. K_in_t = torch.from_numpy(K_in.astype(np.float32)).to(args.device) T_w2c_t = torch.from_numpy(T_world_to_cam.astype(np.float32)).to(args.device) decoded = decode_baseline_motion_tracks( model, rgb_for_model, K_in_t, T_w2c_t, cfg) else: decoded = decode_baseline(model, rgb_for_model, cfg) chunk_buffer.append((abs_step_counter, decoded["raw_out"], None, None)) del chunk_buffer[: max(0, len(chunk_buffer) - args.ensemble_chunks)] if (args.ensemble_chunks > 1 and len(chunk_buffer) > 1 and last_stride > 0): ensembled_raw = _log_pool_raw_out_baseline(last_stride) fake = _FixedRawOutModel(ensembled_raw) decoded = decode_baseline(fake, rgb_for_model, cfg) print(f" ensembled {len(chunk_buffer)} chunks (stride={last_stride}, xyz-mean + logit log-pool)") log_baseline_panels( frame_idx, bgr, joints7, decoded, T_cam_in_rbase, K_in, K_render, Wr, Hr, W, H, img_size, T_world_to_cam) else: # Volume model needs start_pix (current EE projected to image). T_ee_in_rbase = _fk_ee_in_rbase(joints7) ee_world = (T_lfr @ T_ee_in_rbase)[:3, 3] start_pix_xy = _project(ee_world, K_in, T_world_to_cam) start_pix_t = torch.from_numpy(start_pix_xy.astype(np.float32)).unsqueeze(0).to(args.device) # For the hires_ablate model we have to feed past features. At # deploy startup we have no recorded past, so seed it with the # current state replicated past_n times — the model sees "EE has # been stationary at the current position" which is a sensible # prior for the first inference call. past_features = None if kind == "dino_volume_scene_hires_ablate": past_n = int(ck["args"].get("past_n", 0)) if past_n > 0: cur_z = float(ee_world[2]) cur_grip = float(joints7[6]) cur_xyz = ee_world.astype(np.float32) dev = args.device past_features = { "past_z": torch.full((1, past_n), cur_z, dtype=torch.float32, device=dev), "past_grip": torch.full((1, past_n), cur_grip, dtype=torch.float32, device=dev), "past_uv": torch.from_numpy(start_pix_xy.astype(np.float32)) .unsqueeze(0).unsqueeze(0) .expand(1, past_n, 2).to(dev), "past_xyz": torch.from_numpy(cur_xyz) .unsqueeze(0).unsqueeze(0) .expand(1, past_n, 3).to(dev), "current_xyz": torch.from_numpy(cur_xyz) .unsqueeze(0).to(dev), } if is_multiview: # Build per-view (rgb, K_in, T_cam_in_world) lists in the same # order as model.views. Scene cam pose is the locked aruco # extrinsic; wrist cam pose moves with the EE. T_world_ee = T_lfr @ T_ee_in_rbase # (4,4) T_world_wrist_cam = (T_world_ee @ T_ee_wrist_cam if T_ee_wrist_cam is not None else None) bgr_wrist = (_grab_bgr(cam_wrist) if cam_wrist is not None else None) if has_wrist_view and bgr_wrist is None: print(" wrist cam grab failed; retrying"); continue rgb_list = [] K_view_list = [] T_cw_list = [] for v in views: if v == "scene": rgb_list.append(rgb_t) K_view_list.append(K_in) T_cw_list.append(T_cam_in_world) elif v == "wrist": rgb_list.append(_imagenet_norm(bgr_wrist, img_size) .unsqueeze(0).to(args.device)) K_view_list.append(K_in_wrist) T_cw_list.append(T_world_wrist_cam) # start_pix is the EE projected into the PRIMARY view. K_prim = K_view_list[0] T_prim_w2c = np.linalg.inv(T_cw_list[0]) start_pix_xy = _project(ee_world, K_prim, T_prim_w2c) start_pix_t = torch.from_numpy(start_pix_xy.astype(np.float32) ).unsqueeze(0).to(args.device) # Past-trajectory conditioning. On the very first call the # deque is empty → seed with the CURRENT state replicated # past_n times. Same idea the hires_ablate variant used. past_grip_t = past_z_t = None if mv_past_n > 0: cur_grip = float(joints7[6]) cur_z = float(ee_world[2]) while len(past_grip_deque) < mv_past_n: past_grip_deque.append(cur_grip) past_z_deque.append(cur_z) past_grip_t = torch.tensor( past_grip_deque[-mv_past_n:], dtype=torch.float32, device=args.device).unsqueeze(0) past_z_t = torch.tensor( past_z_deque[-mv_past_n:], dtype=torch.float32, device=args.device).unsqueeze(0) if is_v3: decoded = decode_v3( model, rgb_list, cfg, K_view_list, T_cw_list, past_grip=past_grip_t, past_z=past_z_t, border_mask_px=args.border_mask_px) else: decoded = decode_volumetric_multiview( model, rgb_list, start_pix_t, cfg, K_view_list, T_cw_list, border_mask_px=args.border_mask_px, past_grip=past_grip_t, past_z=past_z_t) elif kind == "dino_volume_scene_volumetric": decoded = decode_volumetric(model, rgb_t, start_pix_t, cfg, K_in, T_cam_in_world) else: decoded = decode_outputs(model, rgb_t, start_pix_t, cfg, K_in=K_in, T_cam_in_world=T_cam_in_world, past_features=past_features) # Append the freshly-computed raw_out to the chunk buffer (newest # at the end) and trim to ensemble_chunks. abs_step_counter is the # absolute substep this chunk's prediction starts at. For # multi-view we stash this chunk's per-view K + T_cam_in_world so # _log_pool_volumetric_multiview can reverse-project OLDER chunks' # voxel grids into the NEWEST chunk's frame at log-pool time. if is_multiview: buf_K = list(K_view_list) buf_Tcw = list(T_cw_list) else: buf_K = [K_in] buf_Tcw = [T_cam_in_world] chunk_buffer.append((abs_step_counter, decoded["raw_out"], buf_K, buf_Tcw)) del chunk_buffer[: max(0, len(chunk_buffer) - args.ensemble_chunks)] # Log-pool with prior chunks and re-decode via the stub model. # Stub path bypasses the model forward but reuses workspace + # border masks + argmax + unproject from lib/inference.py:decode. if (args.ensemble_chunks > 1 and len(chunk_buffer) > 1 and last_stride > 0): if is_multiview: ensembled_raw = _log_pool_volumetric_multiview(last_stride) else: ensembled_raw = _log_pool_raw_out(last_stride) fake = _FixedRawOutModel(ensembled_raw) if is_v3: decoded = decode_v3( fake, rgb_list, cfg, K_view_list, T_cw_list, past_grip=past_grip_t, past_z=past_z_t, border_mask_px=args.border_mask_px) elif is_multiview: decoded = decode_volumetric_multiview( fake, rgb_list, start_pix_t, cfg, K_view_list, T_cw_list, border_mask_px=args.border_mask_px, past_grip=past_grip_t, past_z=past_z_t) elif kind == "dino_volume_scene_volumetric": decoded = decode_volumetric(fake, rgb_t, start_pix_t, cfg, K_in, T_cam_in_world) else: decoded = decode_outputs(fake, rgb_t, start_pix_t, cfg, K_in=K_in, T_cam_in_world=T_cam_in_world, past_features=past_features) print(f" ensembled {len(chunk_buffer)} chunks (stride={last_stride})") if is_multiview: bgrs_by_view = {"scene": bgr} K_by_view = {"scene": K_in} T_w2c_by_view = {"scene": np.linalg.inv(T_cam_in_world)} if has_wrist_view and bgr_wrist is not None: bgrs_by_view["wrist"] = bgr_wrist K_by_view["wrist"] = K_in_wrist T_w2c_by_view["wrist"] = np.linalg.inv(T_world_wrist_cam) if is_v3: log_volume_panels_v3( frame_idx, bgr, bgrs_by_view, K_by_view, T_w2c_by_view, joints7, decoded, T_cam_in_rbase, K_render, Wr, Hr, W, H, img_size, views) else: log_volume_panels_multiview( frame_idx, bgr, bgrs_by_view, K_by_view, T_w2c_by_view, joints7, decoded, T_cam_in_rbase, K_render, Wr, Hr, W, H, img_size, views) else: log_volume_panels(frame_idx, bgr, joints7, decoded, T_cam_in_rbase, K_render, Wr, Hr, W, H, img_size) frame_idx += 1 if auto_yes: print("> y (auto — Ctrl-C to pause and re-enter prompt)") key = "y" else: try: key = input("> ").strip().lower() except EOFError: # stdin closed — bail cleanly to home. print() go_home() break # KeyboardInterrupt at the prompt falls through to the # while-loop's outer except: caught there, robot stays put, # next iteration re-prompts. Press `q` to exit. if key == "q": go_home() break elif key == "h": auto_yes = False go_home() rollout_idx += 1 print(f"\n ── rollout {rollout_idx} ──") # New rollout — flush the chunk buffer; old chunks predict a # different absolute timeline. chunk_buffer.clear() last_stride = 0 abs_step_counter = 0 # Drop the past-trajectory deque too — next forward will re-seed # with the (now home) current state, not stale execution history. past_grip_deque.clear() past_z_deque.clear() if recorder is not None: recorder.rotate() bgr_start = _grab_bgr(cam) if bgr_start is not None: recorder.append(bgr_start) if rollout_saver is not None: rollout_saver.rotate() elif key == "c": auto_yes = False print(" re-locking cam pose…") T_cam_in_rbase = pool_exo_aruco( lambda: (_grab_bgr(cam), K_native), min_detections=args.n_detections) T_cam_in_world = to_world(T_cam_in_rbase, T_lfr) T_world_to_cam = np.linalg.inv(T_cam_in_world) print(f" ✓ relocked. cam t = {T_cam_in_rbase[:3, 3]}") elif key == "r": auto_yes = False continue # re-grab + re-infer elif key == "y": try: xyz = decoded["xyz_world"].copy() # (T, 3) quat = decoded["quat"] # (T, 4) grip = decoded["grip"] # (T,) if args.z_offset_m != 0.0: xyz[:, 2] += args.z_offset_m print(f" z_offset {args.z_offset_m*1000:+.1f}mm applied " f"(xyz[:, 2] shifted before IK).") # We only execute the first ``n_substeps`` waypoints before # re-querying, so don't waste IK iterations on the tail. Saves # ~7× wall time when n_substeps=4 vs T=32. (Decoded XYZ for the # FULL window stays in the rerun viz / per-chunk dump — only the # IK call here is shortened.) if args.n_substeps is not None: xyz_ik = xyz[: args.n_substeps] quat_ik = quat[: args.n_substeps] grip_ik = grip[: args.n_substeps] else: xyz_ik, quat_ik, grip_ik = xyz, quat, grip print(f" IK on {len(xyz_ik)} EE poses (seed from current proprio)…") t0 = time.perf_counter() joint_cmds = ee_chain_to_joints(xyz_ik, quat_ik, grip_ik, init_q7=joints7, T_lfr=T_lfr) print(f" IK done in {time.perf_counter() - t0:.2f}s. " f"deploying {len(joint_cmds)} substep(s) via Robot.move_to " f"(first @ {args.first_move_vel_rad_s} rad/s, rest @ " f"{args.max_joint_vel_rad_s} rad/s)") # Persist the decoded chunk (post-ensembling) BEFORE execution. if rollout_saver is not None: # Scene intrinsics + proprio are ALWAYS dumped (any model_kind, # any single/multi-view config). Wrist BGR + pose + K are # included whenever a wrist view was actually consumed by the # model (multiview volume or multiview baseline). save_kwargs: dict = { "K_scene_img_size": K_in, "joints7": joints7, } # Multi-view volume models grab `bgr_wrist`; multi-view baseline # grabs `bgr_wrist_b`. Both feed the same model.views['wrist']. _wrist_bgr_now = None if has_wrist_view: if 'bgr_wrist' in dir() and bgr_wrist is not None: _wrist_bgr_now = bgr_wrist elif 'bgr_wrist_b' in dir() and bgr_wrist_b is not None: _wrist_bgr_now = bgr_wrist_b if _wrist_bgr_now is not None: save_kwargs["bgr_pre_wrist"] = _wrist_bgr_now if T_ee_wrist_cam is not None: T_world_ee_now = T_lfr @ _fk_ee_in_rbase(joints7) save_kwargs["T_wrist_in_world"] = ( T_world_ee_now @ T_ee_wrist_cam) if K_in_wrist is not None: save_kwargs["K_wrist_img_size"] = K_in_wrist rollout_saver.save_chunk(decoded, bgr, T_cam_in_rbase, **save_kwargs) for i, q7 in enumerate(joint_cmds): vel = args.first_move_vel_rad_s if i == 0 else args.max_joint_vel_rad_s target14 = (np.concatenate([np.zeros(7), q7]) if args.arm == "right" else np.concatenate([q7, np.zeros(7)])) info = robot.move_to(target14, vel_rad_s=vel) print(f" [{i+1}/{len(joint_cmds)}] Δ={info.delta_rad:.3f}rad " f"actual {info.actual_s:.2f}s") # Update past-trajectory deque so the next forward sees the # post-substep state. Match the training data: stride=1 per # substep (deploy substeps ≈ recording frames at frame_stride). if mv_past_n > 0: past_grip_deque.append(float(q7[6])) T_ee_after = _fk_ee_in_rbase(q7) past_z_deque.append(float((T_lfr @ T_ee_after)[2, 3])) if len(past_grip_deque) > mv_past_n: past_grip_deque[:] = past_grip_deque[-mv_past_n:] past_z_deque[:] = past_z_deque[-mv_past_n:] if recorder is not None: bgr_after = _grab_bgr(cam) if bgr_after is not None: recorder.append(bgr_after) if rollout_saver is not None: rollout_saver.save_substep_frame(bgr_after, i) if rollout_saver is not None: rollout_saver.end_chunk() # Update buffer bookkeeping: the next forward pass's chunk will # start ``len(joint_cmds)`` substeps after the current one. This # stride is also the offset between consecutive chunks in the # buffer, which the log-pool needs. last_stride = len(joint_cmds) abs_step_counter += len(joint_cmds) auto_yes = True # arm auto-continue for next iter except KeyboardInterrupt: # Ctrl-C mid-chunk: pause robot at last commanded joint # position, fall back to manual prompt. The next forward pass # will re-seed from wherever we actually stopped. auto_yes = False executed = (i + 1) if 'i' in dir() else 0 print(f"\n ⏸ Ctrl-C — paused after {executed}/{len(joint_cmds)} " f"substep(s). Back to manual prompt (y/q/h/r/c).") if rollout_saver is not None: rollout_saver.end_chunk() last_stride = executed abs_step_counter += executed else: auto_yes = False print(f" unknown key '{key}'; press y/q/h/r/c") except KeyboardInterrupt: # Catch-all for Ctrl-C anywhere in the loop — forward pass, # input prompt, or mid-chunk move_to. Robot stays at its last # commanded joints (no go_home); next iteration re-prompts. # Press `q` (not Ctrl-C) to exit the script cleanly. auto_yes = False print("\n ⏸ Ctrl-C — robot paused; back to manual prompt. " "Press 'q' to quit.") if rollout_saver is not None: try: rollout_saver.end_chunk() except Exception: pass continue finally: # Auto-regenerate the eval-report HTML on the lab server so the # rollouts that just landed appear in the table without needing a # manual re-run. SSH via tailscale; fire-and-forget so a slow # tunnel doesn't block the shutdown path. if recorder is not None: try: import subprocess as _sp _sp.Popen( ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=5", "cameronsmith@phe108-yuewang-01", "python /data/cameron/para/robot/yam/code/eval_report.py"], stdout=_sp.DEVNULL, stderr=_sp.DEVNULL, ) print(" triggered eval_report regen on lab (background)") except Exception as _e: print(f" eval_report regen trigger failed: {_e}") # Always smooth back to startup pose before releasing torque, so a # crash mid-session doesn't drop the arm under gravity. go_home() if recorder is not None: recorder.stop() robot.shutdown() cam.close() print("✓ deploy session ended") if __name__ == "__main__": main()