Upload folder using huggingface_hub

.gitattributes

@@ -155,3 +155,10 @@ bundle/artifacts/sim/recorded_motion.gif filter=lfs diff=lfs merge=lfs -text
 bundle/artifacts/sim/scene_motion_fit.png filter=lfs diff=lfs merge=lfs -text
 bundle/artifacts/sim/scene_start_render.png filter=lfs diff=lfs merge=lfs -text
 bundle/artifacts/sim/teddy_mask_overlay.png filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/pi_rollout_rendered_4000_r56.gif filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/pi_rollout_rendered_4000_r56_strip.png filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/pi_rollout_rendered_4000_r96.gif filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/real_vs_scene_start_fitted.png filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/robot_motion_fit_overlay.png filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/scene_background_model.png filter=lfs diff=lfs merge=lfs -text
+bundle/sim_model_based_20260329/artifacts/scene_start_render_fitted.png filter=lfs diff=lfs merge=lfs -text

bundle/sim_model_based_20260329/artifacts/robot_scene_model.json

@@ -0,0 +1,44 @@
+{
+  "session_root": "/workspace/data/teddybear_raw/session_20260327_165944_bear",
+  "fit_rows": [
+    0,
+    52,
+    60,
+    68,
+    76,
+    84,
+    92,
+    100,
+    108,
+    116,
+    124
+  ],
+  "robot_world_rvec": [
+    0.15999206444427902,
+    -0.24953838143443263,
+    -1.8906016307408011
+  ],
+  "robot_world_tvec": [
+    0.4878224859284271,
+    0.48919255742572404,
+    -0.5070228135125122
+  ],
+  "fit_score": 6.463505947380474,
+  "motion_iou_by_row": {
+    "52": 0.06655859635739488,
+    "60": 0.3588341642469731,
+    "68": 0.4452867623651695,
+    "76": 0.48717815284451305,
+    "84": 0.5015739820521984,
+    "92": 0.4417927926329424,
+    "100": 0.39052739334048026,
+    "108": 0.3360236265714658,
+    "116": 0.28526207302709067,
+    "124": 0.2234565091812978
+  },
+  "motion_iou_mean": 0.3536494052619526,
+  "background_path": "/workspace/kinova_scene_sim/outputs/scene_background_model.png",
+  "scene_start_render_path": "/workspace/kinova_scene_sim/outputs/scene_start_render_fitted.png",
+  "real_vs_scene_path": "/workspace/kinova_scene_sim/outputs/real_vs_scene_start_fitted.png",
+  "motion_overlay_path": "/workspace/kinova_scene_sim/outputs/robot_motion_fit_overlay.png"
+}

bundle/sim_model_based_20260329/code/fit_robot_scene_model.py

@@ -0,0 +1,249 @@
+from __future__ import annotations
+
+import argparse
+import json
+import warnings
+from dataclasses import replace
+from pathlib import Path
+
+import cv2
+import numpy as np
+import pandas as pd
+from PIL import Image
+from scipy.optimize import minimize
+
+from scene_tools import (
+    DEFAULT_ROBOT_WORLD_RVEC,
+    DEFAULT_ROBOT_WORLD_TVEC,
+    SceneCalibration,
+    load_scene_calibration,
+    render_robot_mask,
+    render_scene,
+)
+
+
+def parse_rows(spec: str) -> list[int]:
+    return [int(chunk.strip()) for chunk in spec.split(",") if chunk.strip()]
+
+
+def load_rgb_frames(session_root: Path, sync: pd.DataFrame, rows: list[int]) -> dict[int, np.ndarray]:
+    return {row_idx: np.array(Image.open(session_root / sync.iloc[row_idx]["azure_rgb_file"]).convert("RGB")) for row_idx in rows}
+
+
+def load_q_rows(sync: pd.DataFrame, rows: list[int]) -> dict[int, np.ndarray]:
+    return {row_idx: sync.iloc[row_idx][[f"q{i}_deg" for i in range(1, 8)]].to_numpy(np.float64) for row_idx in rows}
+
+
+def compute_motion_mask(base_rgb: np.ndarray, rgb: np.ndarray, roi_mask: np.ndarray, threshold: int = 18) -> np.ndarray:
+    diff = np.abs(rgb.astype(np.int16) - base_rgb.astype(np.int16)).mean(axis=2).astype(np.uint8)
+    mask = ((diff > threshold) & roi_mask).astype(np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((5, 5), dtype=np.uint8))
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((11, 11), dtype=np.uint8))
+    return mask > 0
+
+
+def mask_iou(lhs: np.ndarray, rhs: np.ndarray) -> float:
+    intersection = np.logical_and(lhs, rhs).sum()
+    union = np.logical_or(lhs, rhs).sum()
+    return float(intersection / union) if union else 0.0
+
+
+def fit_robot_transform(
+    calibration: SceneCalibration, rgbs: dict[int, np.ndarray], q_by_row: dict[int, np.ndarray], rows: list[int]
+) -> tuple[np.ndarray, dict[int, np.ndarray], dict[int, float], float]:
+    start_row = rows[0]
+    base_rgb = rgbs[start_row]
+    roi_mask = np.zeros(base_rgb.shape[:2], dtype=bool)
+    roi_mask[250:1079, 850:1919] = True
+    motion_masks = {row_idx: compute_motion_mask(base_rgb, rgbs[row_idx], roi_mask) for row_idx in rows[1:]}
+
+    def score(params: np.ndarray) -> float:
+        robot_world_rvec = params[:3]
+        robot_world_tvec = params[3:]
+        start_mask = render_robot_mask(
+            base_rgb.shape,
+            q_by_row[start_row],
+            calibration.camera_matrix,
+            calibration.rvec,
+            calibration.tvec,
+            robot_world_rvec,
+            robot_world_tvec,
+            extra_dilate=35,
+        ) > 0
+        total = 0.0
+        for row_idx in rows[1:]:
+            current_mask = render_robot_mask(
+                base_rgb.shape,
+                q_by_row[row_idx],
+                calibration.camera_matrix,
+                calibration.rvec,
+                calibration.tvec,
+                robot_world_rvec,
+                robot_world_tvec,
+                extra_dilate=35,
+            ) > 0
+            predicted_change = np.logical_xor(start_mask, current_mask)
+            total += 1.0 - mask_iou(predicted_change, motion_masks[row_idx])
+        return total
+
+    start_params = np.concatenate([DEFAULT_ROBOT_WORLD_RVEC.copy(), DEFAULT_ROBOT_WORLD_TVEC.copy()])
+    result = minimize(score, start_params, method="Powell", options={"maxiter": 40, "disp": False})
+    best_params = np.asarray(result.x, dtype=np.float64)
+
+    start_mask = render_robot_mask(
+        base_rgb.shape,
+        q_by_row[start_row],
+        calibration.camera_matrix,
+        calibration.rvec,
+        calibration.tvec,
+        best_params[:3],
+        best_params[3:],
+        extra_dilate=35,
+    ) > 0
+    ious: dict[int, float] = {}
+    for row_idx in rows[1:]:
+        current_mask = render_robot_mask(
+            base_rgb.shape,
+            q_by_row[row_idx],
+            calibration.camera_matrix,
+            calibration.rvec,
+            calibration.tvec,
+            best_params[:3],
+            best_params[3:],
+            extra_dilate=35,
+        ) > 0
+        ious[row_idx] = mask_iou(np.logical_xor(start_mask, current_mask), motion_masks[row_idx])
+    return best_params, motion_masks, ious, float(result.fun)
+
+
+def build_background(
+    calibration: SceneCalibration,
+    rgbs: dict[int, np.ndarray],
+    q_by_row: dict[int, np.ndarray],
+    rows: list[int],
+    robot_world_rvec: np.ndarray,
+    robot_world_tvec: np.ndarray,
+) -> np.ndarray:
+    base_rgb = rgbs[rows[0]]
+    masked_frames = []
+    valid_masks = []
+    for row_idx in rows:
+        robot_mask = render_robot_mask(
+            base_rgb.shape,
+            q_by_row[row_idx],
+            calibration.camera_matrix,
+            calibration.rvec,
+            calibration.tvec,
+            robot_world_rvec,
+            robot_world_tvec,
+            extra_dilate=65,
+        ) > 0
+        masked_frames.append(rgbs[row_idx].astype(np.float32))
+        valid_masks.append(~robot_mask)
+
+    frame_stack = np.stack(masked_frames, axis=0)
+    valid_stack = np.stack(valid_masks, axis=0)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        background = np.zeros_like(base_rgb, dtype=np.float32)
+        for channel in range(3):
+            channel_values = np.where(valid_stack, frame_stack[..., channel], np.nan)
+            background[..., channel] = np.nanmedian(channel_values, axis=0)
+    hole_mask = np.isnan(background[..., 0])
+    background = np.nan_to_num(background, nan=0.0).astype(np.uint8)
+    if np.any(hole_mask):
+        background = cv2.cvtColor(
+            cv2.inpaint(cv2.cvtColor(background, cv2.COLOR_RGB2BGR), hole_mask.astype(np.uint8) * 255, 9, cv2.INPAINT_TELEA),
+            cv2.COLOR_BGR2RGB,
+        )
+    return background
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--session-root", default="/workspace/data/teddybear_raw/session_20260327_165944_bear")
+    parser.add_argument("--fit-rows", default="0,52,60,68,76,84,92,100,108,116,124")
+    parser.add_argument("--output-dir", default="/workspace/kinova_scene_sim/outputs")
+    args = parser.parse_args()
+
+    session_root = Path(args.session_root)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    sync = pd.read_csv(session_root / "sync_index.csv")
+    fit_rows = parse_rows(args.fit_rows)
+    calibration = load_scene_calibration(session_root, sync_row_index=0)
+    rgbs = load_rgb_frames(session_root, sync, fit_rows)
+    q_by_row = load_q_rows(sync, fit_rows)
+
+    best_params, motion_masks, ious, fit_score = fit_robot_transform(calibration, rgbs, q_by_row, fit_rows)
+    calibration = replace(
+        calibration,
+        robot_world_rvec=best_params[:3].copy(),
+        robot_world_tvec=best_params[3:].copy(),
+    )
+    background = build_background(calibration, rgbs, q_by_row, fit_rows, best_params[:3], best_params[3:])
+
+    start_row = fit_rows[0]
+    representative_row = fit_rows[min(len(fit_rows) - 1, 4)]
+    start_q = q_by_row[start_row]
+    start_rgb = rgbs[start_row]
+    start_render = render_scene(calibration, start_q, background=background.copy(), alpha=1.0)
+
+    start_mask = render_robot_mask(
+        start_rgb.shape,
+        start_q,
+        calibration.camera_matrix,
+        calibration.rvec,
+        calibration.tvec,
+        calibration.robot_world_rvec,
+        calibration.robot_world_tvec,
+    ) > 0
+    representative_mask = render_robot_mask(
+        start_rgb.shape,
+        q_by_row[representative_row],
+        calibration.camera_matrix,
+        calibration.rvec,
+        calibration.tvec,
+        calibration.robot_world_rvec,
+        calibration.robot_world_tvec,
+    ) > 0
+    predicted_change = np.logical_xor(start_mask, representative_mask)
+    motion_overlay = rgbs[representative_row].copy()
+    motion_overlay[motion_masks[representative_row]] = np.array([0, 255, 0], dtype=np.uint8)
+    motion_overlay[predicted_change] = np.array([255, 0, 0], dtype=np.uint8)
+
+    background_path = output_dir / "scene_background_model.png"
+    start_render_path = output_dir / "scene_start_render_fitted.png"
+    real_vs_scene_path = output_dir / "real_vs_scene_start_fitted.png"
+    motion_overlay_path = output_dir / "robot_motion_fit_overlay.png"
+    scene_model_json = output_dir / "robot_scene_model.json"
+
+    Image.fromarray(background).save(background_path)
+    Image.fromarray(start_render).save(start_render_path)
+    Image.fromarray(np.concatenate([start_rgb, start_render], axis=1)).save(real_vs_scene_path)
+    Image.fromarray(motion_overlay).save(motion_overlay_path)
+
+    payload = {
+        "session_root": str(session_root),
+        "fit_rows": fit_rows,
+        "robot_world_rvec": calibration.robot_world_rvec.tolist(),
+        "robot_world_tvec": calibration.robot_world_tvec.tolist(),
+        "fit_score": fit_score,
+        "motion_iou_by_row": {str(key): value for key, value in ious.items()},
+        "motion_iou_mean": float(np.mean(list(ious.values()))),
+        "background_path": str(background_path),
+        "scene_start_render_path": str(start_render_path),
+        "real_vs_scene_path": str(real_vs_scene_path),
+        "motion_overlay_path": str(motion_overlay_path),
+    }
+    scene_model_json.write_text(json.dumps(payload, indent=2))
+
+    print(f"saved scene model to {scene_model_json}")
+    print(f"robot_world_rvec={calibration.robot_world_rvec.tolist()}")
+    print(f"robot_world_tvec={calibration.robot_world_tvec.tolist()}")
+    print(f"motion_iou_mean={payload['motion_iou_mean']:.4f}")
+
+
+if __name__ == "__main__":
+    main()

bundle/sim_model_based_20260329/code/render_pose_sequence.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import cv2
+import numpy as np
+import pandas as pd
+from PIL import Image
+from scipy.optimize import least_squares
+
+from scene_tools import create_background_inpaint, kinova_fk_points_and_tool_pose, load_scene_calibration, render_scene
+
+
+def euler_xyz_deg_to_matrix(theta_x_deg: float, theta_y_deg: float, theta_z_deg: float) -> np.ndarray:
+    rx, ry, rz = np.deg2rad([theta_x_deg, theta_y_deg, theta_z_deg])
+    cx, sx = np.cos(rx), np.sin(rx)
+    cy, sy = np.cos(ry), np.sin(ry)
+    cz, sz = np.cos(rz), np.sin(rz)
+    rot_x = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]], dtype=np.float64)
+    rot_y = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]], dtype=np.float64)
+    rot_z = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]], dtype=np.float64)
+    return rot_z @ rot_y @ rot_x
+
+
+def solve_ik(target_pose: np.ndarray, q_init_deg: np.ndarray) -> np.ndarray:
+    target_pos = target_pose[:3].astype(np.float64)
+    target_rot = euler_xyz_deg_to_matrix(*target_pose[3:6])
+    q_init = np.asarray(q_init_deg, dtype=np.float64)
+
+    def residual(q_deg: np.ndarray) -> np.ndarray:
+        _, tool = kinova_fk_points_and_tool_pose(q_deg)
+        pos_err = tool[:3, 3] - target_pos
+        rot_err_mat = tool[:3, :3].T @ target_rot
+        rot_err, _ = cv2.Rodrigues(rot_err_mat)
+        regularize = 0.01 * (q_deg - q_init)
+        return np.concatenate([12.0 * pos_err, 0.75 * rot_err.reshape(-1), regularize])
+
+    result = least_squares(residual, q_init, method="lm", max_nfev=120)
+    return result.x.astype(np.float64)
+
+
+def load_target_poses(args: argparse.Namespace) -> tuple[np.ndarray, np.ndarray]:
+    if args.poses_json:
+        payload = json.loads(Path(args.poses_json).read_text())
+        if isinstance(payload, dict) and "actions" in payload:
+            poses = np.asarray(payload["actions"], dtype=np.float64)
+        else:
+            poses = np.asarray(payload, dtype=np.float64)
+        start_q = np.asarray(args.q_init_deg, dtype=np.float64)
+        return poses, start_q
+
+    sync = pd.read_csv(Path(args.session_root) / "sync_index.csv")
+    rows = sync.iloc[args.sync_start : args.sync_stop : args.sync_step]
+    poses = rows[
+        ["tool_x_m", "tool_y_m", "tool_z_m", "tool_theta_x_deg", "tool_theta_y_deg", "tool_theta_z_deg", "gripper_pos"]
+    ].to_numpy(np.float64)
+    start_q = rows.iloc[0][[f"q{i}_deg" for i in range(1, 8)]].to_numpy(np.float64)
+    return poses, start_q
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--session-root", default="/workspace/data/teddybear_raw/session_20260327_165944_bear")
+    parser.add_argument("--output", default="/workspace/kinova_scene_sim/outputs/rollout.gif")
+    parser.add_argument("--poses-json", default=None, help="JSON file containing a list of 7D tool poses.")
+    parser.add_argument(
+        "--q-init-deg",
+        nargs=7,
+        type=float,
+        default=[0.0008167303, 340.0004272461, 180.0014038086, 213.9954223633, 359.9988708496, 309.9989318848, 90.0014266968],
+        help="Initial joint seed for IK when using --poses-json.",
+    )
+    parser.add_argument("--sync-start", type=int, default=0)
+    parser.add_argument("--sync-stop", type=int, default=130)
+    parser.add_argument("--sync-step", type=int, default=8)
+    parser.add_argument("--fps", type=int, default=5)
+    args = parser.parse_args()
+
+    calibration = load_scene_calibration(args.session_root, sync_row_index=0)
+    background = create_background_inpaint(calibration.rgb)
+    poses, q_seed = load_target_poses(args)
+
+    frames: list[Image.Image] = []
+    q_current = q_seed.copy()
+    for idx, pose in enumerate(poses):
+        q_current = solve_ik(pose, q_current)
+        frame = render_scene(calibration, q_current, background=background)
+        # Mark the target tool center on the rendered frame.
+        points_world, _ = kinova_fk_points_and_tool_pose(q_current)
+        _ = points_world  # quiet lint-like readers; pose solving already uses the FK.
+        frames.append(Image.fromarray(frame))
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    frames[0].save(
+        output,
+        save_all=True,
+        append_images=frames[1:],
+        loop=0,
+        duration=int(1000 / max(args.fps, 1)),
+    )
+    print(f"saved rollout gif to {output}")
+
+
+if __name__ == "__main__":
+    main()

bundle/sim_model_based_20260329/code/rollout_pi_policy_rendered.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from dataclasses import replace
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+from render_pose_sequence import solve_ik
+from scene_tools import load_scene_calibration, render_scene, robot_tool_position_world
+
+
+def main() -> None:
+    t0 = time.time()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint-dir", required=True)
+    parser.add_argument("--session-root", default="/workspace/data/teddybear_raw/session_20260327_165944_bear")
+    parser.add_argument("--sync-row-index", type=int, default=52)
+    parser.add_argument("--config-name", default="pi05_kinova_teddybear")
+    parser.add_argument("--prompt", default="pick up the teddy bear and place it in the red box")
+    parser.add_argument("--rollout-steps", type=int, default=32)
+    parser.add_argument("--scene-model-json", default="/workspace/kinova_scene_sim/outputs/robot_scene_model.json")
+    parser.add_argument("--output-gif", default=None)
+    parser.add_argument("--output-json", default="/workspace/kinova_scene_sim/outputs/pi_rollout_rendered.json")
+    args = parser.parse_args()
+
+    session_root = Path(args.session_root)
+    sync = pd.read_csv(session_root / "sync_index.csv")
+    row = sync.iloc[args.sync_row_index]
+    calibration = load_scene_calibration(session_root, sync_row_index=0)
+    if args.scene_model_json and Path(args.scene_model_json).exists():
+        scene_model = json.loads(Path(args.scene_model_json).read_text())
+        calibration = replace(
+            calibration,
+            robot_world_rvec=np.asarray(scene_model["robot_world_rvec"], dtype=np.float64),
+            robot_world_tvec=np.asarray(scene_model["robot_world_tvec"], dtype=np.float64),
+        )
+        background_path = Path(scene_model["background_path"])
+    else:
+        background_path = session_root / row["azure_rgb_file"]
+    base_background = np.array(Image.open(background_path).convert("RGB"))
+    wrist_rgb = np.array(Image.open(session_root / row["rgb_file"]).convert("RGB"))
+
+    current_state = np.array(
+        [
+            row["tool_x_m"],
+            row["tool_y_m"],
+            row["tool_z_m"],
+            row["tool_theta_x_deg"],
+            row["tool_theta_y_deg"],
+            row["tool_theta_z_deg"],
+            row["gripper_pos"],
+        ],
+        dtype=np.float32,
+    )
+    q_current = row[[f"q{i}_deg" for i in range(1, 8)]].to_numpy(np.float64)
+    print(f"[rendered] start row={args.sync_row_index}", flush=True)
+    print("[rendered] importing openpi modules", flush=True)
+    from openpi.policies import policy_config as _policy_config
+    from openpi.training import config as _config
+
+    print(f"[rendered] imported openpi in {time.time() - t0:.2f}s", flush=True)
+    train_config = _config.get_config(args.config_name)
+    policy = _policy_config.create_trained_policy(train_config, Path(args.checkpoint_dir))
+    print(f"[rendered] policy ready in {time.time() - t0:.2f}s", flush=True)
+
+    states = [current_state.copy()]
+    qs = [q_current.copy()]
+    chunks = []
+    rendered_files: list[str] = []
+    teddy_distances_m = [
+        float(
+            np.linalg.norm(
+                robot_tool_position_world(
+                    current_state[:3], calibration.robot_world_rvec, calibration.robot_world_tvec
+                )[:2]
+                - calibration.teddy_world_center[:2]
+            )
+        )
+    ]
+
+    output_dir = Path(args.output_json).parent / "rendered_frames"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for step_idx in range(args.rollout_steps):
+        rendered = render_scene(calibration, q_current, background=base_background.copy(), alpha=1.0)
+        rendered_path = output_dir / f"frame_{step_idx:03d}.png"
+        Image.fromarray(rendered).save(rendered_path)
+        rendered_files.append(str(rendered_path))
+
+        results = policy.infer(
+            {
+                "observation/state": current_state,
+                "observation/image": rendered,
+                "observation/wrist_image": wrist_rgb,
+                "prompt": args.prompt,
+            }
+        )
+        chunk = np.asarray(results["actions"], dtype=np.float32)
+        chunks.append(chunk)
+        current_state = chunk[0].copy()
+        q_current = solve_ik(current_state, q_current)
+        states.append(current_state.copy())
+        qs.append(q_current.copy())
+        teddy_distances_m.append(
+            float(
+                np.linalg.norm(
+                    robot_tool_position_world(
+                        current_state[:3], calibration.robot_world_rvec, calibration.robot_world_tvec
+                    )[:2]
+                    - calibration.teddy_world_center[:2]
+                )
+            )
+        )
+        print(
+            f"[rendered] step={step_idx + 1}/{args.rollout_steps} "
+            f"xyz={current_state[:3].tolist()} grip={float(current_state[6]):.3f} "
+            f"teddy_xy_m={teddy_distances_m[-1]:.4f}",
+            flush=True,
+        )
+
+    output = {
+        "checkpoint_dir": str(Path(args.checkpoint_dir)),
+        "session_root": str(session_root),
+        "sync_row_index": int(args.sync_row_index),
+        "states": np.stack(states, axis=0).tolist(),
+        "actions": np.stack(states[1:], axis=0).tolist(),
+        "chunks": np.stack(chunks, axis=0).tolist(),
+        "joint_sequences_deg": np.stack(qs, axis=0).tolist(),
+        "rendered_files": rendered_files,
+        "teddy_xy_distance_m": teddy_distances_m,
+    }
+    output_path = Path(args.output_json)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(json.dumps(output, indent=2))
+    output_gif = Path(args.output_gif) if args.output_gif else output_path.with_suffix(".gif")
+    frames = [Image.open(frame_path) for frame_path in rendered_files]
+    frames[0].save(
+        output_gif,
+        save_all=True,
+        append_images=frames[1:],
+        loop=0,
+        duration=130,
+    )
+    print(f"[rendered] saved rollout to {output_path}", flush=True)
+    print(f"[rendered] saved gif to {output_gif}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

bundle/sim_model_based_20260329/code/scene_tools.py

@@ -0,0 +1,487 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+import pandas as pd
+from PIL import Image
+
+
+# Manually fitted table corners for the synced overhead frame used for episode 0.
+# Order: top-left, top-right, bottom-right, bottom-left.
+TABLE_CORNERS_PX = np.array(
+    [
+        [192.0, 138.0],
+        [1635.0, 46.0],
+        [1712.0, 787.0],
+        [185.0, 858.0],
+    ],
+    dtype=np.float64,
+)
+
+# Assumed tabletop size for the large brown work surface in the overhead image.
+TABLE_SIZE_M = (1.52, 0.76)
+
+# Tuned against the pre-grasp motion silhouette.
+CAMERA_FOCAL_PX = 650.0
+ROBOT_BASE_WORLD = np.array([0.50, 0.45, -0.45], dtype=np.float64)
+ROBOT_BASE_YAW_RAD = -2.0
+DEFAULT_ROBOT_WORLD_RVEC = np.array([0.06924932, -0.1712883, -1.95079235], dtype=np.float64)
+DEFAULT_ROBOT_WORLD_TVEC = np.array([0.50851769, 0.49565435, -0.50654742], dtype=np.float64)
+
+# Kinova tool-center offset relative to the interface frame from FK.
+TOOL_OFFSET_M = 0.12
+
+
+@dataclass(frozen=True)
+class SceneCalibration:
+    session_root: Path
+    sync_row_index: int
+    azure_rgb_seq: int
+    azure_depth_seq: int
+    robot_seq: int
+    rgb_path: Path
+    depth_path: Path
+    rgb: np.ndarray
+    depth: np.ndarray
+    sync_row: dict[str, Any]
+    camera_matrix: np.ndarray
+    rvec: np.ndarray
+    tvec: np.ndarray
+    table_corners_px: np.ndarray
+    table_size_m: tuple[float, float]
+    robot_base_world: np.ndarray
+    robot_base_yaw_rad: float
+    robot_world_rvec: np.ndarray
+    robot_world_tvec: np.ndarray
+    box_mask: np.ndarray
+    teddy_mask: np.ndarray
+    box_world_polygon: np.ndarray
+    teddy_world_center: np.ndarray
+    teddy_world_extent: np.ndarray
+    table_depth_mm: float
+    box_height_m: float
+    teddy_height_m: float
+
+
+def _to_rgb(path: Path) -> np.ndarray:
+    return np.array(Image.open(path).convert("RGB"))
+
+
+def _to_depth(path: Path) -> np.ndarray:
+    return np.array(Image.open(path))
+
+
+def load_sync_dataframe(session_root: Path) -> pd.DataFrame:
+    return pd.read_csv(session_root / "sync_index.csv")
+
+
+def load_scene_calibration(session_root: str | Path, sync_row_index: int = 0) -> SceneCalibration:
+    session_root = Path(session_root)
+    sync = load_sync_dataframe(session_root)
+    row = sync.iloc[sync_row_index]
+
+    azure_rgb_seq = int(row["azure_rgb_seq"])
+    azure_depth_seq = int(row["azure_depth_seq"])
+    robot_seq = int(row["robot_seq"])
+
+    rgb_path = session_root / row["azure_rgb_file"]
+    depth_path = session_root / row["azure_depth_file"]
+    rgb = _to_rgb(rgb_path)
+    depth = _to_depth(depth_path)
+
+    camera_matrix, rvec, tvec = solve_table_camera(TABLE_CORNERS_PX, TABLE_SIZE_M, CAMERA_FOCAL_PX)
+    box_mask = detect_red_box(rgb)
+    teddy_mask = detect_teddy(rgb)
+    box_world_polygon = image_mask_to_world_polygon(box_mask, camera_matrix, rvec, tvec)
+    teddy_world_center, teddy_world_extent = mask_centroid_and_extent_world(teddy_mask, camera_matrix, rvec, tvec)
+
+    table_depth_mm = estimate_table_depth_mm(depth)
+    box_height_m = estimate_height_from_depth(rgb, depth, box_mask, table_depth_mm)
+    teddy_height_m = estimate_height_from_depth(rgb, depth, teddy_mask, table_depth_mm)
+
+    return SceneCalibration(
+        session_root=session_root,
+        sync_row_index=sync_row_index,
+        azure_rgb_seq=azure_rgb_seq,
+        azure_depth_seq=azure_depth_seq,
+        robot_seq=robot_seq,
+        rgb_path=rgb_path,
+        depth_path=depth_path,
+        rgb=rgb,
+        depth=depth,
+        sync_row=row.to_dict(),
+        camera_matrix=camera_matrix,
+        rvec=rvec,
+        tvec=tvec,
+        table_corners_px=TABLE_CORNERS_PX.copy(),
+        table_size_m=TABLE_SIZE_M,
+        robot_base_world=ROBOT_BASE_WORLD.copy(),
+        robot_base_yaw_rad=float(ROBOT_BASE_YAW_RAD),
+        robot_world_rvec=DEFAULT_ROBOT_WORLD_RVEC.copy(),
+        robot_world_tvec=DEFAULT_ROBOT_WORLD_TVEC.copy(),
+        box_mask=box_mask,
+        teddy_mask=teddy_mask,
+        box_world_polygon=box_world_polygon,
+        teddy_world_center=teddy_world_center,
+        teddy_world_extent=teddy_world_extent,
+        table_depth_mm=table_depth_mm,
+        box_height_m=box_height_m,
+        teddy_height_m=teddy_height_m,
+    )
+
+
+def solve_table_camera(
+    table_corners_px: np.ndarray, table_size_m: tuple[float, float], focal_px: float
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    width_m, height_m = table_size_m
+    object_points = np.array(
+        [
+            [-width_m / 2, -height_m / 2, 0.0],
+            [width_m / 2, -height_m / 2, 0.0],
+            [width_m / 2, height_m / 2, 0.0],
+            [-width_m / 2, height_m / 2, 0.0],
+        ],
+        dtype=np.float64,
+    )
+    camera_matrix = np.array(
+        [[focal_px, 0.0, 960.0], [0.0, focal_px, 540.0], [0.0, 0.0, 1.0]],
+        dtype=np.float64,
+    )
+    ok, rvec, tvec = cv2.solvePnP(object_points, table_corners_px, camera_matrix, None, flags=cv2.SOLVEPNP_IPPE)
+    if not ok:
+        raise RuntimeError("solvePnP failed for table camera fit")
+    return camera_matrix, rvec, tvec
+
+
+def detect_red_box(rgb: np.ndarray) -> np.ndarray:
+    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
+    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
+    mask = (((h < 12) | (h > 170)) & (s > 60) & (v > 80)).astype(np.uint8)
+    num, labels, stats, centers = cv2.connectedComponentsWithStats(mask, connectivity=8)
+    best = None
+    best_area = -1
+    for i in range(1, num):
+        x, y, w, h_box, area = stats[i]
+        cx, cy = centers[i]
+        if area > 50_000 and x > 900 and y < 700 and area > best_area:
+            best_area = int(area)
+            best = labels == i
+    if best is None:
+        raise RuntimeError("failed to detect red box component")
+    return best
+
+
+def detect_teddy(rgb: np.ndarray) -> np.ndarray:
+    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
+    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
+    mask = ((h > 5) & (h < 22) & (s > 15) & (s < 120) & (v > 110) & (v < 245)).astype(np.uint8)
+    num, labels, stats, centers = cv2.connectedComponentsWithStats(mask, connectivity=8)
+    best = None
+    best_area = -1
+    for i in range(1, num):
+        x, y, w, h_box, area = stats[i]
+        cx, cy = centers[i]
+        if 5_000 < area < 50_000 and 850 < cx < 1250 and 350 < cy < 700 and area > best_area:
+            best_area = int(area)
+            best = labels == i
+    if best is None:
+        raise RuntimeError("failed to detect teddy component")
+    return best
+
+
+def estimate_table_depth_mm(depth: np.ndarray) -> float:
+    center_region = depth[150:430, 120:520]
+    valid = center_region[center_region > 0]
+    if len(valid) == 0:
+        raise RuntimeError("no valid table depth values")
+    return float(np.percentile(valid, 50))
+
+
+def estimate_height_from_depth(rgb: np.ndarray, depth: np.ndarray, rgb_mask: np.ndarray, table_depth_mm: float) -> float:
+    ys, xs = np.where(rgb_mask)
+    if len(xs) == 0:
+        return 0.0
+    xd = np.clip((xs * depth.shape[1] / rgb.shape[1]).astype(int), 0, depth.shape[1] - 1)
+    yd = np.clip((ys * depth.shape[0] / rgb.shape[0]).astype(int), 0, depth.shape[0] - 1)
+    values = depth[yd, xd]
+    values = values[values > 0]
+    if len(values) == 0:
+        return 0.0
+    # Use a low percentile to reduce the effect of table pixels around the object.
+    object_depth_mm = float(np.percentile(values, 10))
+    return max(0.0, (table_depth_mm - object_depth_mm) / 1000.0)
+
+
+def image_mask_to_world_polygon(mask: np.ndarray, camera_matrix: np.ndarray, rvec: np.ndarray, tvec: np.ndarray) -> np.ndarray:
+    points = np.column_stack(np.where(mask > 0))[:, ::-1].astype(np.float64)
+    rect = cv2.minAreaRect(points.astype(np.float32))
+    corners = cv2.boxPoints(rect).astype(np.float64)
+    return pixels_to_world_on_table(corners, camera_matrix, rvec, tvec)
+
+
+def mask_centroid_and_extent_world(
+    mask: np.ndarray, camera_matrix: np.ndarray, rvec: np.ndarray, tvec: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    points = np.column_stack(np.where(mask > 0))[:, ::-1].astype(np.float64)
+    world_points = pixels_to_world_on_table(points, camera_matrix, rvec, tvec)
+    center = world_points.mean(axis=0)
+    extent = world_points.max(axis=0) - world_points.min(axis=0)
+    return center, extent
+
+
+def pixels_to_world_on_table(
+    pixels_uv: np.ndarray, camera_matrix: np.ndarray, rvec: np.ndarray, tvec: np.ndarray
+) -> np.ndarray:
+    rotation, _ = cv2.Rodrigues(rvec)
+    camera_position = (-rotation.T @ tvec).reshape(3)
+    inv_camera = np.linalg.inv(camera_matrix)
+    world_points = []
+    for u, v in pixels_uv:
+        ray_cam = inv_camera @ np.array([u, v, 1.0], dtype=np.float64)
+        ray_cam /= np.linalg.norm(ray_cam)
+        ray_world = rotation.T @ ray_cam
+        if abs(ray_world[2]) < 1e-8:
+            continue
+        scale = -camera_position[2] / ray_world[2]
+        world_points.append(camera_position + scale * ray_world)
+    if not world_points:
+        raise RuntimeError("failed to back-project pixels to the table plane")
+    return np.stack(world_points, axis=0)
+
+
+def create_background_inpaint(rgb: np.ndarray) -> np.ndarray:
+    mask = np.zeros(rgb.shape[:2], dtype=np.uint8)
+    polygon = np.array(
+        [[1080, 520], [1919, 520], [1919, 1079], [980, 1079], [980, 850], [1080, 760]],
+        dtype=np.int32,
+    )
+    cv2.fillPoly(mask, [polygon], 255)
+    cv2.rectangle(mask, (1120, 300), (1450, 680), 0, -1)
+    cv2.rectangle(mask, (860, 350), (1180, 700), 0, -1)
+    return cv2.cvtColor(cv2.inpaint(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR), mask, 7, cv2.INPAINT_TELEA), cv2.COLOR_BGR2RGB)
+
+
+def kinova_fk_points_world(q_deg: np.ndarray, base_world: np.ndarray, base_yaw_rad: float) -> np.ndarray:
+    points, _ = kinova_fk_points_and_tool_pose(q_deg)
+    cy, sy = np.cos(base_yaw_rad), np.sin(base_yaw_rad)
+    base_rotation = np.array([[cy, -sy, 0.0], [sy, cy, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64)
+    return points @ base_rotation.T + base_world.reshape(1, 3)
+
+
+def transform_robot_points(points_robot: np.ndarray, robot_world_rvec: np.ndarray, robot_world_tvec: np.ndarray) -> np.ndarray:
+    rotation, _ = cv2.Rodrigues(np.asarray(robot_world_rvec, dtype=np.float64).reshape(3, 1))
+    translation = np.asarray(robot_world_tvec, dtype=np.float64).reshape(1, 3)
+    return np.asarray(points_robot, dtype=np.float64) @ rotation.T + translation
+
+
+def robot_tool_position_world(
+    tool_position_robot: np.ndarray, robot_world_rvec: np.ndarray, robot_world_tvec: np.ndarray
+) -> np.ndarray:
+    return transform_robot_points(np.asarray(tool_position_robot, dtype=np.float64).reshape(1, 3), robot_world_rvec, robot_world_tvec)[0]
+
+
+def kinova_fk_points_and_tool_pose(q_deg: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    q = np.deg2rad(np.asarray(q_deg, dtype=np.float64))
+    params = [
+        (np.pi, 0.0, 0.0, 0.0),
+        (np.pi / 2, 0.0, -(0.1564 + 0.1284), q[0]),
+        (np.pi / 2, 0.0, -(0.0054 + 0.0064), q[1] + np.pi),
+        (np.pi / 2, 0.0, -(0.2104 + 0.2104), q[2] + np.pi),
+        (np.pi / 2, 0.0, -(0.0064 + 0.0064), q[3] + np.pi),
+        (np.pi / 2, 0.0, -(0.2084 + 0.1059), q[4] + np.pi),
+        (np.pi / 2, 0.0, 0.0, q[5] + np.pi),
+        (np.pi, 0.0, -(0.1059 + 0.0615), q[6] + np.pi),
+    ]
+
+    def dh(alpha: float, a: float, d: float, theta: float) -> np.ndarray:
+        ca, sa = np.cos(alpha), np.sin(alpha)
+        ct, st = np.cos(theta), np.sin(theta)
+        return np.array(
+            [
+                [ct, -st * ca, st * sa, a * ct],
+                [st, ct * ca, -ct * sa, a * st],
+                [0.0, sa, ca, d],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            dtype=np.float64,
+        )
+
+    transform = np.eye(4, dtype=np.float64)
+    points = []
+    for alpha, a, d, theta in params:
+        transform = transform @ dh(alpha, a, d, theta)
+        points.append(transform[:3, 3].copy())
+
+    tool_transform = transform.copy()
+    tool_transform[:3, 3] = tool_transform[:3, 3] + tool_transform[:3, 2] * TOOL_OFFSET_M
+    points.append(tool_transform[:3, 3].copy())
+    return np.stack(points, axis=0), tool_transform
+
+
+def project_world_points(points_world: np.ndarray, camera_matrix: np.ndarray, rvec: np.ndarray, tvec: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    image_points, _ = cv2.projectPoints(points_world.astype(np.float64), rvec, tvec, camera_matrix, None)
+    rotation, _ = cv2.Rodrigues(rvec)
+    camera_points = (rotation @ points_world.T + tvec).T
+    return image_points[:, 0, :], camera_points[:, 2]
+
+
+def draw_robot(
+    image: np.ndarray,
+    q_deg: np.ndarray,
+    camera_matrix: np.ndarray,
+    rvec: np.ndarray,
+    tvec: np.ndarray,
+    base_world: np.ndarray | None = None,
+    base_yaw_rad: float | None = None,
+    robot_world_rvec: np.ndarray | None = None,
+    robot_world_tvec: np.ndarray | None = None,
+    color: tuple[int, int, int] = (232, 242, 250),
+    alpha: float = 1.0,
+) -> np.ndarray:
+    canvas = image.copy()
+    if robot_world_rvec is not None and robot_world_tvec is not None:
+        points_robot, _ = kinova_fk_points_and_tool_pose(q_deg)
+        points_world = transform_robot_points(points_robot, robot_world_rvec, robot_world_tvec)
+    else:
+        if base_world is None or base_yaw_rad is None:
+            raise ValueError("base_world/base_yaw_rad or robot_world_rvec/robot_world_tvec must be provided")
+        points_world = kinova_fk_points_world(q_deg, base_world, base_yaw_rad)
+    uv, z = project_world_points(points_world, camera_matrix, rvec, tvec)
+
+    overlay = canvas.copy()
+    link_palette = [
+        (214, 224, 232),
+        (222, 232, 239),
+        (214, 224, 232),
+        (206, 220, 229),
+        (198, 214, 224),
+        (190, 208, 220),
+        (182, 202, 215),
+        (172, 196, 212),
+    ]
+    outline = (70, 86, 98)
+    joint_fill = (230, 236, 240)
+
+    for idx, (p0, p1, depth_value) in enumerate(zip(uv[:-1], uv[1:], z[:-1], strict=False)):
+        base_thickness = max(14, int(52 / max(depth_value, 0.35)))
+        fill_color = link_palette[min(idx, len(link_palette) - 1)]
+        cv2.line(
+            overlay,
+            tuple(np.round(p0).astype(int)),
+            tuple(np.round(p1).astype(int)),
+            outline,
+            base_thickness + 8,
+            lineType=cv2.LINE_AA,
+        )
+        cv2.line(
+            overlay,
+            tuple(np.round(p0).astype(int)),
+            tuple(np.round(p1).astype(int)),
+            fill_color,
+            base_thickness,
+            lineType=cv2.LINE_AA,
+        )
+    for point, depth_value in zip(uv, z, strict=False):
+        radius = max(10, int(24 / max(depth_value, 0.35)))
+        center = tuple(np.round(point).astype(int))
+        cv2.circle(overlay, center, radius + 4, outline, -1, lineType=cv2.LINE_AA)
+        cv2.circle(overlay, center, radius, joint_fill, -1, lineType=cv2.LINE_AA)
+    if alpha >= 1.0:
+        return overlay
+    return cv2.addWeighted(overlay, alpha, canvas, 1.0 - alpha, 0.0)
+
+
+def render_robot_mask(
+    image_shape: tuple[int, int] | tuple[int, int, int],
+    q_deg: np.ndarray,
+    camera_matrix: np.ndarray,
+    rvec: np.ndarray,
+    tvec: np.ndarray,
+    robot_world_rvec: np.ndarray,
+    robot_world_tvec: np.ndarray,
+    extra_dilate: int = 0,
+) -> np.ndarray:
+    height, width = image_shape[:2]
+    points_robot, _ = kinova_fk_points_and_tool_pose(q_deg)
+    points_world = transform_robot_points(points_robot, robot_world_rvec, robot_world_tvec)
+    uv, z = project_world_points(points_world, camera_matrix, rvec, tvec)
+    mask = np.zeros((height, width), dtype=np.uint8)
+    for p0, p1, depth_value in zip(uv[:-1], uv[1:], z[:-1], strict=False):
+        thickness = max(18, int(60 / max(depth_value, 0.35)))
+        cv2.line(
+            mask,
+            tuple(np.round(p0).astype(int)),
+            tuple(np.round(p1).astype(int)),
+            255,
+            thickness,
+            lineType=cv2.LINE_AA,
+        )
+    for point, depth_value in zip(uv, z, strict=False):
+        radius = max(12, int(26 / max(depth_value, 0.35)))
+        cv2.circle(mask, tuple(np.round(point).astype(int)), radius, 255, -1, lineType=cv2.LINE_AA)
+    if extra_dilate > 0:
+        kernel = np.ones((extra_dilate, extra_dilate), dtype=np.uint8)
+        mask = cv2.dilate(mask, kernel, iterations=1)
+    return mask
+
+
+def render_scene(
+    calibration: SceneCalibration,
+    q_deg: np.ndarray,
+    background: np.ndarray | None = None,
+    color: tuple[int, int, int] = (232, 242, 250),
+    alpha: float = 1.0,
+) -> np.ndarray:
+    if background is None:
+        background = create_background_inpaint(calibration.rgb)
+    return draw_robot(
+        background,
+        q_deg,
+        calibration.camera_matrix,
+        calibration.rvec,
+        calibration.tvec,
+        base_world=calibration.robot_base_world,
+        base_yaw_rad=calibration.robot_base_yaw_rad,
+        robot_world_rvec=calibration.robot_world_rvec,
+        robot_world_tvec=calibration.robot_world_tvec,
+        color=color,
+        alpha=alpha,
+    )
+
+
+def scene_to_jsonable(calibration: SceneCalibration) -> dict[str, Any]:
+    return {
+        "session_root": str(calibration.session_root),
+        "sync_row_index": calibration.sync_row_index,
+        "azure_rgb_seq": calibration.azure_rgb_seq,
+        "azure_depth_seq": calibration.azure_depth_seq,
+        "robot_seq": calibration.robot_seq,
+        "rgb_path": str(calibration.rgb_path),
+        "depth_path": str(calibration.depth_path),
+        "table_corners_px": calibration.table_corners_px.tolist(),
+        "table_size_m": list(calibration.table_size_m),
+        "camera_focal_px": float(calibration.camera_matrix[0, 0]),
+        "rvec": calibration.rvec.reshape(-1).tolist(),
+        "tvec": calibration.tvec.reshape(-1).tolist(),
+        "robot_base_world": calibration.robot_base_world.tolist(),
+        "robot_base_yaw_rad": float(calibration.robot_base_yaw_rad),
+        "robot_world_rvec": calibration.robot_world_rvec.tolist(),
+        "robot_world_tvec": calibration.robot_world_tvec.tolist(),
+        "sync_row": calibration.sync_row,
+        "table_depth_mm": float(calibration.table_depth_mm),
+        "box_height_m": float(calibration.box_height_m),
+        "teddy_height_m": float(calibration.teddy_height_m),
+        "box_world_polygon": calibration.box_world_polygon.tolist(),
+        "teddy_world_center": calibration.teddy_world_center.tolist(),
+        "teddy_world_extent": calibration.teddy_world_extent.tolist(),
+    }
+
+
+def save_scene_json(calibration: SceneCalibration, path: str | Path) -> None:
+    path = Path(path)
+    path.write_text(json.dumps(scene_to_jsonable(calibration), indent=2))