"""V4HEAD-VIDEO VARIANT (yams_any4d, 2026-07-11): minimal diff from deploy_with_action_chunking.py. The v4 model forward is replaced by a file-bridge round-trip to v4head_deploy_server.py (a4d-bw env): live scene frame -> video-model GENERATION (cond on the single live frame) -> cached V4Head -> 41 (xyz, quat, grip). Generated 41-frame video is logged to a gen/video rerun panel. Camera / ArUco lock / y-r-h-c-q keys / IK / robot execution are untouched. Original doc follows. 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, to_jpeg 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", default="", help="UNUSED in the v4head-video variant (model lives in the bridge server)") p.add_argument("--bridge_dir", default=str(Path.home() / "any4d_work/deploy_bridge"), help="Shared dir with v4head_deploy_server.py") p.add_argument("--bridge_timeout_s", type=float, default=300.0) 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" args.device = device # NOTE(yams_any4d): no local model. Inference = file bridge to # v4head_deploy_server.py. cfg keeps only what downstream geometry uses. model = None cfg = {"img_size": 504, "n_window": 41} kind = "v4head_video" is_baseline = is_baseline_mt = is_baseline_diff = is_baseline_mv = False is_v3 = is_v4 = is_multiview = False views = ["scene"] has_scene_view = True has_wrist_view = False mv_past_n = 0 past_grip_deque: list[float] = [] past_z_deque: list[float] = [] bridge = Path(args.bridge_dir) bridge.mkdir(parents=True, exist_ok=True) _req_counter = int(time.time()) % 100000 def _bridge_infer(bgr_frame, joints7_arr, T_ee_rbase): """Send frame+state, block until the server responds. Returns (decoded dict, gen_frames list[bgr]) or (None, None) on timeout/error.""" nonlocal _req_counter _req_counter += 1 rid = f"{_req_counter:06d}" ok, jbuf = cv2.imencode(".jpg", bgr_frame, [cv2.IMWRITE_JPEG_QUALITY, 92]) tmp = bridge / f"request_{rid}.tmp.npz" np.savez(tmp, bgr_jpeg=np.frombuffer(jbuf.tobytes(), dtype=np.uint8), joints7=joints7_arr.astype(np.float32), T_ee_in_rbase=T_ee_rbase.astype(np.float64)) (bridge / f"request_{rid}.tmp.npz").rename(bridge / f"request_{rid}.npz") resp_path = bridge / f"response_{rid}.npz" t0 = time.time() print(f" → bridge request {rid} (video generation ~60-120s)…", flush=True) while not resp_path.exists(): if time.time() - t0 > args.bridge_timeout_s: print(" ✗ bridge timeout"); return None, None time.sleep(0.5) time.sleep(0.3) resp = np.load(resp_path, allow_pickle=True) status = str(resp["status"]) resp_path.unlink(missing_ok=True) if status != "ok": print(f" ✗ bridge error: {status}"); return None, None gen_frames = [cv2.imdecode(j, cv2.IMREAD_COLOR) for j in resp["gen_jpegs"]] panels_l = {k: bytes(resp[k]) for k in ("panel_pca_tok", "panel_pca_up", "panel_heatmap", "panel_grip", "panel_rot") if k in resp.files} decoded_l = {"xyz_world": resp["xyz"].astype(np.float64), "quat": resp["quat"].astype(np.float64), "grip": resp["grip"].astype(np.float64), "raw_out": {}, "panels": panels_l} print(f" ← bridge response in {time.time()-t0:.1f}s " f"({len(gen_frames)} gen frames)") return decoded_l, gen_frames print(f"v4head-video deploy: bridge dir = {bridge} (start " f"v4head_deploy_server.py in a4d-bw first)") 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, send_default_blueprint=False) try: import rerun.blueprint as rrb rr.send_blueprint(rrb.Blueprint(rrb.Grid( rrb.Spatial2DView(name="scene RGB", origin="scene/rgb"), rrb.Spatial2DView(name="scene composite", origin="scene/composite"), rrb.Spatial2DView(name="GEN video", origin="gen/video"), rrb.Spatial2DView(name="pred kp", origin="pred/kp"), rrb.Spatial2DView(name="heatmap grid", origin="pred/heatmap_grid"), rrb.Spatial2DView(name="PCA tokens", origin="pred/pca_scene"), rrb.Spatial2DView(name="PCA upsampled", origin="pred/pca_upsampled"), rrb.Spatial2DView(name="grip grid", origin="pred/grip_grid"), rrb.Spatial2DView(name="rot grid", origin="pred/rot_grid"), rrb.TextDocumentView(name="status", origin="status"), ), collapse_panels=True)) print(" ✓ v4head-video blueprint sent") except Exception as _be: print(f" blueprint send failed ({_be})") 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) # NOTE(yams_any4d): model forward replaced by the bridge round-trip. T_ee_in_rbase = _fk_ee_in_rbase(joints7) ee_world = (T_lfr @ T_ee_in_rbase)[:3, 3] decoded, gen_frames = _bridge_infer(bgr, joints7, T_ee_in_rbase) if decoded is None: print(" bridge failed; press r to retry"); decoded = None rr.set_time("step", sequence=frame_idx) rr.log("scene/rgb", to_jpeg(cv2.resize(bgr, (Wr, Hr)), target_w=None)) if decoded is not None: # generated video on its own frame timeline for gt_i, gf in enumerate(gen_frames or []): rr.set_time("gen_frame", sequence=frame_idx * 64 + gt_i) rr.log("gen/video", to_jpeg(gf, target_w=None)) rr.set_time("step", sequence=frame_idx) # predicted 41-step EE trajectory + GT-current marker on the live frame overlay = cv2.resize(bgr, (Wr, Hr)).copy() T_w2c_live = np.linalg.inv(T_cam_in_world) pts = [] for t_i in range(decoded["xyz_world"].shape[0]): uv = _project(decoded["xyz_world"][t_i], K_render, T_w2c_live) pts.append(uv) pts_i = np.round(np.asarray(pts)).astype(int) for t_i in range(len(pts_i) - 1): cv2.line(overlay, tuple(pts_i[t_i]), tuple(pts_i[t_i + 1]), (255, 255, 255), 1, cv2.LINE_AA) for t_i in range(len(pts_i)): hue = int(170 * t_i / max(len(pts_i) - 1, 1)) col = cv2.cvtColor(np.uint8([[[hue, 255, 255]]]), cv2.COLOR_HSV2BGR)[0, 0] cv2.circle(overlay, tuple(pts_i[t_i]), 4, tuple(int(c) for c in col), -1, cv2.LINE_AA) cur_uv = np.round(_project(ee_world, K_render, T_w2c_live)).astype(int) cv2.circle(overlay, tuple(cur_uv), 9, (255, 255, 255), 2, cv2.LINE_AA) rr.log("pred/kp", to_jpeg(overlay, target_w=None)) _pmap = {"panel_pca_tok": "pred/pca_scene", "panel_pca_up": "pred/pca_upsampled", "panel_heatmap": "pred/heatmap_grid", "panel_grip": "pred/grip_grid", "panel_rot": "pred/rot_grid"} for (_pk, _dest) in _pmap.items(): _jb = decoded.get("panels", {}).get(_pk) if _jb: rr.log(_dest, rr.EncodedImage(contents=_jb, media_type="image/jpeg")) rr.log("status", rr.TextDocument( f"plan @ step {frame_idx}\nxyz[0]={decoded['xyz_world'][0].round(3)}" f" xyz[-1]={decoded['xyz_world'][-1].round(3)}\n" f"grip[0]={decoded['grip'][0]:.2f} frames={len(gen_frames or [])}")) print(f" pred: xyz[0]={decoded['xyz_world'][0].round(3)} " f"xyz[-1]={decoded['xyz_world'][-1].round(3)} " f"grip[0]={decoded['grip'][0]:.2f}") 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: if decoded is None: print(" no valid prediction (bridge failed) — press r to retry") continue 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()