scripts/render_conditioning_set.py

#!/usr/bin/env python
"""Export per-frame geometry-conditioning channels for a video diffusion model.

Given a planner-produced kinematic trajectory (robot joints + gripper +
per-object world poses) for a pick-and-place, this renders the
ControlNet-depth / Cosmos-Transfer-style channels a conditioned video model
consumes: RGB, 16-bit depth, an instance segmentation, and per-frame pose
metadata, plus the camera and a human-viewable preview clip.

Input contract (produced by a planner elsewhere in the pipeline, fixed):
a ``.npz`` with ``timestamps`` (T,), ``joint_positions`` (T, 7), ``gripper``
(T,) in DROID's 0=open convention, ``object_poses`` (T, n_obj, 4, 4) world
poses, ``object_names`` (n_obj,). Every field but ``timestamps`` is treated
as optional here: a missing/malformed field disables exactly the channel(s)
it feeds (no robot rendered without joint_positions, no objects without
object_poses) with a loud warning, rather than a KeyError mid-loop -- the
planner is produced by a different, concurrently-running process, so this
script cannot assume its output is perfect on the first try.

The static background (table, shelf, whatever the camera actually saw) comes
from one real depth capture of a DROID episode, unprojected into a coloured
point cloud and then *carved* of the robot/object silhouettes the camera saw
at that instant -- exactly render_pointcloud_scene.py's approach, reused
directly rather than reimplemented (see the imports below). Skipping the
carve would render the arm and object twice: once frozen in the capture
photo, once moving under the planner's trajectory -- the same surface
represented twice, precisely the wrong signal for geometry conditioning.

Per-frame instance masks (background / robot / object 1 / object 2 / ...) use
the same two-pass depth-agreement trick as
:class:`fpgm.objects.scene.SceneRenderer` (src/fpgm/objects/scene.py:184-230),
generalised from one object to N: render everything together for the
occlusion-correct image, then render each entity alone (everything else
detached) and credit it with whichever pixels its own unoccluded depth
agrees with the combined depth. This needs no alpha/object-id channel from
OSMesa (which this build does not reliably expose either).

Usage:
    PYTHONPATH=src python scripts/render_conditioning_set.py \\
        --npz outputs/planner_trajectory.npz \\
        --episode AUTOLab+0d4edc83+2023-10-21-19h-07m-04s --camera ext1 --clip 15:26 \\
        --object-asset brick=outputs/.../objects_sam3d_multi/15_26 \\
        --out-dir outputs/conditioning_set --mode fixed
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import resource
import sys
import time
from dataclasses import dataclass, field
from pathlib import Path

import h5py
import numpy as np
import trimesh
from PIL import Image

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "src"))

# NOTE on ordering (also present, and also flagged by ruff's isort check, in
# render_pointcloud_scene.py -- this repo's existing convention): `pyrender`
# is imported last, after `fpgm.robot.render`, whose import sets
# PYOPENGL_PLATFORM=osmesa and stubs pyrender.viewer as a side effect. Both
# must happen before pyrender/OpenGL is first imported anywhere in the
# process (see that module's docstring) -- ruff's import sorter does not know
# about this constraint and will suggest hoisting `pyrender` above the
# alphabetised block; do not apply that fix, it silently breaks headless GL.
from fpgm.data.droid_raw import cam2world_vector_to_world2cam  # noqa: E402
from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.geometry.transforms import invert_se3  # noqa: E402
from fpgm.objects.scene import free_camera  # noqa: E402
from fpgm.robot.render import RobotRenderer  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.types import DataError  # noqa: E402
from fpgm.utils.io import ensure_dir  # noqa: E402
from fpgm.utils.logging import get_logger, setup_logging  # noqa: E402
from fpgm.viz.overlays import VideoWriter  # noqa: E402

import pyrender  # noqa: E402

# Reuse render_pointcloud_scene.py's point-cloud unprojection/carving/drawing
# machinery wholesale rather than reimplementing it -- same importlib trick
# that module itself uses for run_object_pipeline.py (sys.modules[...] MUST be
# set before exec_module, or that module's own @dataclass decorators break).
_rps_spec = importlib.util.spec_from_file_location(
    "render_pointcloud_scene", REPO_ROOT / "scripts" / "render_pointcloud_scene.py"
)
rps = importlib.util.module_from_spec(_rps_spec)
sys.modules[_rps_spec.name] = rps
_rps_spec.loader.exec_module(rps)
rop = rps.rop  # run_object_pipeline, already imported the same way by rps itself

logger = get_logger("render_conditioning_set")

_DEFAULT_URDF = rop._DEFAULT_URDF

#: OSMesa (software GL) leaks ~180MB/frame on this host regardless of scene
#: reuse (measured in scripts/render_robot_random.py); recreating the renderer
#: every N frames bounds that growth instead of OOMing the shared host over a
#: long export. Same mitigation, same default interval.
_RENDERER_RECYCLE = 10

#: Generic small-object stand-in used only when no real mesh asset is given
#: for an object name -- keeps the export running (never a crash) but must
#: never be mistaken for the real reconstructed geometry.
_PLACEHOLDER_EXTENT_M = 0.08
_PLACEHOLDER_PALETTE_RGB: tuple[tuple[int, int, int], ...] = (
    (220, 120, 60), (60, 160, 220), (160, 220, 60), (220, 60, 160), (60, 220, 160),
)

#: Franka/Robotiq links whose FK positions define the TCP, matching the
#: convention already established by fpgm.robot.kinematics.ArmKinematics (see
#: its _build_flange_table docstring): fingertip midpoint for position,
#: gripper-base rotation for orientation. Not reused via that class directly
#: since it additionally demands a validated all-Z-axis arm for its fast
#: Jacobian path -- more machinery than one FK lookup per frame needs here.
_TCP_FINGER_LINKS = ("left_inner_finger", "right_inner_finger")
_TCP_ROTATION_LINK = "robotiq_85_base_link"

_BASIS_POINTS = np.array(
    [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    p.add_argument("--npz", type=Path, required=True, help="planner trajectory .npz, see docstring")
    p.add_argument("--action-json", type=Path, default=None, help="sidecar action-spec JSON")
    p.add_argument("--episode", required=True, help="DROID episode the background capture is from")
    p.add_argument("--camera", default="ext1", choices=["ext1", "ext2"])
    p.add_argument(
        "--clip", required=True,
        help='PointWorld clip key whose initial_depth/rgb is the background, e.g. "15:26"',
    )
    p.add_argument(
        "--object-asset", action="append", default=None, metavar="NAME=DIR",
        help="stage_root directory (containing objects/2_mesh and objects/3_align, i.e. an "
        "outputs/<episode>/objects* tree) providing the real mesh + capture-time alignment "
        "for object NAME. Repeatable. An object present in --npz with no matching entry "
        "here renders as a placeholder box and is not carved out of the background cloud.",
    )
    p.add_argument("--urdf", type=Path, default=_DEFAULT_URDF)
    p.add_argument("--out-dir", type=Path, required=True)
    p.add_argument(
        "--mode", choices=["fixed", "orbit"], default="fixed",
        help="fixed = the real capture viewpoint; orbit = swing the camera to reveal 3D structure",
    )
    p.add_argument("--orbit-degrees", type=float, default=50.0)
    p.add_argument("--width", type=int, default=640)
    p.add_argument("--height", type=int, default=360)
    p.add_argument("--max-depth", type=float, default=3.0, help="drop far cloud points, metres")
    p.add_argument("--point-radius", type=int, default=2)
    p.add_argument(
        "--robot-dilate-px", type=int, default=2,
        help="dilate the robot silhouette before carving, to catch depth bleed at edges",
    )
    p.add_argument("--object-pad-m", type=float, default=0.02)
    p.add_argument(
        "--renderer-recycle", type=int, default=_RENDERER_RECYCLE,
        help="recreate the OSMesa renderer every N frames (0 disables; see module docstring)",
    )
    p.add_argument("--max-frames", type=int, default=None, help="render at most N frames (debug)")
    return p.parse_args()


# --------------------------------------------------------------------------- #
# Planner .npz: load field-by-field, never crash mid-loop on a missing one.
# --------------------------------------------------------------------------- #


@dataclass
class PlannerTrajectory:
    """The subset of the input contract that actually loaded, plus what didn't.

    Every array field is ``None`` if the source ``.npz`` did not have it (or
    had it in the wrong shape) -- callers switch off exactly the channel that
    field feeds rather than assuming it is present.
    """

    timestamps: np.ndarray  # (T,) float64, seconds
    joint_positions: np.ndarray | None  # (T, 7) or None
    gripper: np.ndarray  # (T,), defaults to all-open (0.0) if absent
    object_poses: np.ndarray | None  # (T, n_obj, 4, 4) or None
    object_names: list[str]  # len 0 if object_poses is None
    fps: float
    warnings: list[str] = field(default_factory=list)

    @property
    def n_frames(self) -> int:
        return int(self.timestamps.shape[0])


_EXPECTED_FPS = 15.0


def _load_planner_npz(path: Path) -> PlannerTrajectory:
    """Load the planner's ``.npz`` against the fixed input contract.

    ``timestamps`` is the only field treated as truly required (without it
    there is no way to know how many frames to render at all); everything
    else degrades to "channel disabled" plus a logged warning, per the module
    docstring.
    """
    if not path.exists():
        raise DataError(f"planner npz not found: {path}")

    warns: list[str] = []
    with np.load(path, allow_pickle=False) as npz:
        keys = set(npz.files)
        if "timestamps" not in keys:
            raise DataError(f"{path}: missing required 'timestamps' field; frame count unknown")
        timestamps = np.asarray(npz["timestamps"], dtype=np.float64).reshape(-1)
        t = timestamps.shape[0]
        if t == 0:
            raise DataError(f"{path}: 'timestamps' is empty")

        dt = np.diff(timestamps)
        fps = _EXPECTED_FPS
        if dt.size:
            if not np.allclose(dt, dt[0], atol=1e-6):
                warns.append(
                    f"timestamps not uniformly spaced (dt {dt.min():.4f}-{dt.max():.4f}s)"
                )
            if dt[0] > 0:
                fps = 1.0 / float(dt[0])
            if abs(fps - _EXPECTED_FPS) > 0.5:
                warns.append(
                    f"fps inferred from timestamps ({fps:.2f}) is far from expected {_EXPECTED_FPS}"
                )

        joint_positions = None
        if "joint_positions" not in keys:
            warns.append("no 'joint_positions' field: the robot will NOT be rendered")
        else:
            jp = np.asarray(npz["joint_positions"], dtype=np.float64)
            if jp.shape != (t, 7):
                warns.append(
                    f"'joint_positions' has shape {jp.shape}, expected {(t, 7)}: "
                    "the robot will NOT be rendered"
                )
            else:
                joint_positions = jp

        gripper = np.zeros(t, dtype=np.float64)
        if "gripper" not in keys:
            warns.append("no 'gripper' field: defaulting every frame to open (0.0)")
        else:
            g = np.asarray(npz["gripper"], dtype=np.float64).reshape(-1)
            if g.shape[0] != t:
                warns.append(f"'gripper' has length {g.shape[0]}, expected {t}: defaulting to open")
            else:
                gripper = g

        object_names: list[str] = []
        if "object_names" in keys:
            object_names = [str(n) for n in np.asarray(npz["object_names"]).reshape(-1)]

        object_poses = None
        if "object_poses" not in keys:
            warns.append("no 'object_poses' field: no objects will be rendered")
            object_names = []
        else:
            op = np.asarray(npz["object_poses"], dtype=np.float64)
            if op.ndim != 4 or op.shape[0] != t or op.shape[2:] != (4, 4):
                warns.append(
                    f"'object_poses' has shape {op.shape}, expected {(t, '<n_obj>', 4, 4)}: "
                    "no objects will be rendered"
                )
                object_names = []
            elif op.shape[1] != len(object_names):
                warns.append(
                    f"'object_poses' carries {op.shape[1]} objects but 'object_names' has "
                    f"{len(object_names)} entries: no objects will be rendered"
                )
                object_names = []
            else:
                object_poses = op

    for w in warns:
        logger.warning("%s: %s", path, w)

    return PlannerTrajectory(
        timestamps=timestamps, joint_positions=joint_positions, gripper=gripper,
        object_poses=object_poses, object_names=object_names, fps=fps, warnings=warns,
    )


# --------------------------------------------------------------------------- #
# Object assets: real mesh + capture-time pose (for carving), or a graceful
# placeholder when no real asset is available for a name.
# --------------------------------------------------------------------------- #


@dataclass
class ObjectAsset:
    name: str
    mesh: trimesh.Trimesh
    #: mesh-canonical -> world transform at the instant the background depth
    #: was captured -- needed only to carve this object's real-capture ghost
    #: out of the point cloud. ``None`` when unknown (placeholder, or no
    #: --object-asset given), in which case carving that object is skipped.
    capture_transform: np.ndarray | None
    is_real: bool


def _parse_object_assets(pairs: list[str] | None) -> dict[str, Path]:
    result: dict[str, Path] = {}
    for pair in pairs or []:
        if "=" not in pair:
            raise SystemExit(f"--object-asset must be NAME=DIR, got {pair!r}")
        name, _, dir_str = pair.partition("=")
        result[name] = Path(dir_str)
    return result


def _placeholder_object_mesh(color_rgb: tuple[int, int, int]) -> trimesh.Trimesh:
    tri = trimesh.creation.box(extents=[_PLACEHOLDER_EXTENT_M] * 3)
    fill = np.array([*color_rgb, 255], dtype=np.uint8)
    tri.visual.vertex_colors = np.tile(fill, (len(tri.vertices), 1))
    return tri


def _load_object_assets(names: list[str], asset_dirs: dict[str, Path]) -> list[ObjectAsset]:
    assets: list[ObjectAsset] = []
    for i, name in enumerate(names):
        color = _PLACEHOLDER_PALETTE_RGB[i % len(_PLACEHOLDER_PALETTE_RGB)]
        stage_root = asset_dirs.get(name)
        if stage_root is None:
            logger.warning(
                "no --object-asset given for %r; substituting a %.0fcm placeholder box "
                "(NOT the real geometry) and skipping cloud carving for it",
                name, _PLACEHOLDER_EXTENT_M * 100,
            )
            assets.append(ObjectAsset(name, _placeholder_object_mesh(color), None, False))
            continue
        try:
            tri, transform = rps._load_object(stage_root)
        except Exception as exc:  # noqa: BLE001 - degrade to a placeholder, never crash the export
            logger.warning(
                "failed to load object asset %r from %s (%s); substituting a placeholder box",
                name, stage_root, exc,
            )
            assets.append(ObjectAsset(name, _placeholder_object_mesh(color), None, False))
            continue
        assets.append(ObjectAsset(name, tri, transform, True))
    return assets


# --------------------------------------------------------------------------- #
# Rendering: robot + N objects, one pyrender scene, per-entity instance masks.
# --------------------------------------------------------------------------- #


class _ConditioningRenderer:
    """Robot + N independently-posed objects, with a disjoint per-entity mask.

    Generalises :class:`fpgm.objects.scene.SceneRenderer`'s two-pass
    depth-agreement trick (src/fpgm/objects/scene.py:184-230) from exactly one
    object to however many the trajectory carries: render everything together
    for the occlusion-correct composite, then for each object render it alone
    (every other node -- every robot link, every other object -- detached) and
    credit it with whichever combined-mask pixels its own unoccluded depth
    agrees with (within :attr:`_MASK_DEPTH_ATOL`). Every combined-mask pixel no
    object claims is the robot's, by exclusion -- exact at silhouette
    boundaries, and needs no object-id/alpha channel from OSMesa (which this
    build does not reliably expose either, see RobotRenderer's own note).

    Costs one extra render pass per object per frame -- cheap next to scene
    construction, which is what actually dominates OSMesa's software runtime.
    """

    _MASK_DEPTH_ATOL = 1e-4

    def __init__(
        self,
        link_meshes: dict[str, list[tuple[trimesh.Trimesh, np.ndarray]]],
        object_meshes: list[trimesh.Trimesh],
        width: int,
        height: int,
        **robot_kwargs: object,
    ) -> None:
        self._robot = RobotRenderer(link_meshes, width, height, **robot_kwargs)
        self._object_nodes: list[pyrender.Node] = []
        for tri in object_meshes:
            pmesh = pyrender.Mesh.from_trimesh(tri, smooth=False)
            node = pyrender.Node(mesh=pmesh, matrix=np.eye(4))
            self._robot.scene.add_node(node)
            self._object_nodes.append(node)
        # Flat list of every robot link's node, toggled off/on per object-alone
        # pass below -- reaches into RobotRenderer's own bookkeeping rather
        # than duplicating it, same integration point SceneRenderer uses.
        self._link_nodes = [
            node for entries in self._robot._link_nodes.values() for node, _ in entries
        ]

    def render(
        self,
        link_poses: dict[str, np.ndarray],
        object_poses: list[np.ndarray],
        camera: Camera,
    ):
        """Returns ``(combined_result, robot_mask, object_masks)``.

        ``object_masks[i]`` corresponds to the i-th mesh this renderer was
        constructed with, in order.
        """
        for node, pose in zip(self._object_nodes, object_poses, strict=True):
            self._robot.scene.set_pose(node, np.asarray(pose, dtype=np.float64))
        combined = self._robot.render(link_poses, camera)

        object_masks: list[np.ndarray] = []
        for i, _node in enumerate(self._object_nodes):
            others = self._link_nodes + [n for j, n in enumerate(self._object_nodes) if j != i]
            for n in others:
                self._robot.scene.remove_node(n)
            try:
                _, solo_depth = self._robot._renderer.render(self._robot.scene)
            finally:
                for n in others:
                    self._robot.scene.add_node(n)
            solo_depth = np.asarray(solo_depth, dtype=np.float32)
            object_masks.append(
                combined.mask
                & (solo_depth > 0.0)
                & (np.abs(solo_depth - combined.depth) <= self._MASK_DEPTH_ATOL)
            )

        any_object = np.zeros(combined.mask.shape, dtype=bool)
        for m in object_masks:
            any_object |= m
        robot_mask = combined.mask & ~any_object
        return combined, robot_mask, object_masks

    def close(self) -> None:
        self._robot.close()


def _composite_depth(
    mesh_depth: np.ndarray, uv: np.ndarray, point_depth: np.ndarray, radius: int
) -> np.ndarray:
    """Depth channel, metres; 0 = no data.

    Mirrors the compositing rule the main loop's rgb canvas actually ends up
    with, not a naive nearer-of-the-two z-test: after
    ``render_pointcloud_scene._draw_points`` stamps cloud points (respecting
    its own visibility test), the caller unconditionally overwrites every
    mesh-covered pixel with the mesh colour (``canvas[combined.mask] = ...``)
    -- so the mesh silhouette always wins on-screen regardless of whether a
    background point happens to be nearer there. A two-way z-test here would
    let this channel disagree with what rgb/seg actually show at exactly
    those pixels (checked against real data: it does happen, wherever the
    background pokes into the mesh's 2D silhouette without truly being in
    front of it in the mesh's own occluded corner). So: mesh depth is
    authoritative on its own footprint, and only pixels with no mesh at all
    (``mesh_depth <= 0``, matching ``seg == 0``) are filled from the carved
    cloud -- nearest cloud point among those that land there, same stamp
    radius as ``_draw_points``.
    """
    h, w = mesh_depth.shape
    out = mesh_depth.copy()
    empty = out <= 0.0

    u = np.round(uv[:, 0]).astype(np.int64)
    v = np.round(uv[:, 1]).astype(np.int64)
    keep = (point_depth > 0) & (u >= 0) & (u < w) & (v >= 0) & (v < h)
    u, v, point_depth = u[keep], v[keep], point_depth[keep]

    order = np.argsort(point_depth)[::-1]  # far first, so nearer points overwrite
    u, v, point_depth = u[order], v[order], point_depth[order]

    r = int(radius)
    for du in range(-r, r + 1):
        for dv in range(-r, r + 1):
            if du * du + dv * dv > r * r:
                continue
            uu = np.clip(u + du, 0, w - 1)
            vv = np.clip(v + dv, 0, h - 1)
            can_write = empty[vv, uu]
            out[vv[can_write], uu[can_write]] = point_depth[can_write]
    return out


def _tcp_pose(link_poses: dict[str, np.ndarray]) -> np.ndarray | None:
    """base_link -> TCP: fingertip-midpoint position, gripper-base orientation.

    Matches the convention :class:`fpgm.robot.kinematics.ArmKinematics` is
    built against (see its ``_build_flange_table`` docstring) without pulling
    in that class, which additionally validates an all-Z-axis arm -- more than
    one FK dict lookup per frame needs.
    """
    try:
        left = link_poses[_TCP_FINGER_LINKS[0]][:3, 3]
        right = link_poses[_TCP_FINGER_LINKS[1]][:3, 3]
        rot = link_poses[_TCP_ROTATION_LINK][:3, :3]
    except KeyError as exc:
        logger.warning("URDF has no link %s; cannot compute tcp_pose", exc)
        return None
    pose = np.eye(4, dtype=np.float64)
    pose[:3, :3] = rot
    pose[:3, 3] = (left + right) / 2.0
    return pose


def _camera_to_world(camera: Camera) -> np.ndarray:
    """Recover `camera`'s (4, 4) cam->world matrix via its public point-transform API.

    :class:`~fpgm.geometry.camera.Camera` deliberately exposes only
    point-transform methods, not a raw extrinsic matrix -- this mirrors the
    exact same reconstruction ``fpgm.robot.render._camera_to_world_matrix``
    already uses for that reason, rather than reaching into a private
    attribute.
    """
    transformed = camera.cam_to_world(_BASIS_POINTS)
    origin = transformed[0]
    rotation = (transformed[1:] - origin).T
    m = np.eye(4, dtype=np.float64)
    m[:3, :3] = rotation
    m[:3, 3] = origin
    return m


def _world_to_cam_matrix(camera: Camera) -> np.ndarray:
    return invert_se3(_camera_to_world(camera))


def _rss_mb() -> int:
    """Process peak RSS so far, MB. Linux ``ru_maxrss`` is kilobytes and monotonic
    non-decreasing across the process's lifetime -- exactly what is needed to
    confirm the OSMesa-leak mitigation is bounding growth over a long export."""
    return int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024)


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #


def main() -> int:
    args = parse_args()
    setup_logging()

    traj = _load_planner_npz(args.npz)
    n_frames = traj.n_frames
    if args.max_frames is not None:
        n_frames = min(n_frames, args.max_frames)
        logger.info("capping render to %d/%d frames (--max-frames)", n_frames, traj.n_frames)

    out_dir = args.out_dir
    rgb_dir = ensure_dir(out_dir / "rgb")
    depth_dir = ensure_dir(out_dir / "depth")
    seg_dir = ensure_dir(out_dir / "seg")
    pose_dir = ensure_dir(out_dir / "pose")

    # --- Background: one real depth capture, carved of whatever robot/object
    # the camera actually saw frozen in that frame (mandatory, see module docstring).
    ctx = rop._load_clip_context(args.episode, args.camera, args.clip)
    metadata = rop._load_metadata(args.episode)
    serial = metadata[f"{args.camera}_cam_serial"]
    depth_u16 = rop._load_initial_depth(args.episode, args.clip, serial)
    cloud_xyz, cloud_rgb = rps._scene_point_cloud(ctx, depth_u16, args.max_depth)
    logger.info("background cloud: %d points before carving", cloud_xyz.shape[0])

    robot = RobotModel(str(args.urdf), load_meshes=True)

    # The trajectory row / camera pose the depth frame was actually captured
    # at -- from the REAL episode, independent of the planner's (possibly
    # synthetic) trajectory, since carving must remove what the camera really
    # saw, not what the plan says happened.
    capture_row = int(round(ctx.timing.clip_frame_to_trajectory_frame(0)))
    with h5py.File(rop._load_trajectory_h5(args.episode), "r") as f:
        real_joint_positions = np.asarray(f["observation/robot_state/joint_positions"])
        real_gripper = np.asarray(f["observation/robot_state/gripper_position"])
        extrinsics_6d = np.asarray(f[f"observation/camera_extrinsics/{serial}_left"])
    capture_row = int(
        np.clip(capture_row, 0, min(real_joint_positions.shape[0], extrinsics_6d.shape[0]) - 1)
    )
    capture_w2c = cam2world_vector_to_world2cam(extrinsics_6d[capture_row])

    cloud_xyz, cloud_rgb = rps._drop_robot_points(
        cloud_xyz, cloud_rgb, robot,
        real_joint_positions[capture_row], float(real_gripper[capture_row]),
        ctx, capture_w2c, args.robot_dilate_px,
    )

    asset_dirs = _parse_object_assets(args.object_asset)
    object_names = traj.object_names
    object_assets = _load_object_assets(object_names, asset_dirs)
    for asset in object_assets:
        if asset.capture_transform is not None:
            cloud_xyz, cloud_rgb = rps._drop_object_points(
                cloud_xyz, cloud_rgb, asset.mesh, asset.capture_transform,
                args.object_pad_m, asset.name,
            )
        else:
            logger.warning(
                "object %r has no known capture-time pose (no real --object-asset alignment); "
                "it cannot be carved out of the background cloud, so a ghost of it may remain",
                asset.name,
            )
    logger.info("background cloud: %d points after carving", cloud_xyz.shape[0])

    K = ctx.camera_annot.K.scaled(args.width, args.height)
    base_cam = Camera(K, capture_w2c)
    cloud_centre = cloud_xyz.mean(axis=0) if cloud_xyz.shape[0] else np.zeros(3)
    cam_origin = base_cam.cam_to_world(np.zeros((1, 3)))[0]
    orbit_distance = float(np.linalg.norm(cam_origin - cloud_centre))

    link_meshes = robot.visual_meshes() if traj.joint_positions is not None else {}
    object_meshes = [asset.mesh for asset in object_assets]

    def _new_renderer() -> _ConditioningRenderer:
        return _ConditioningRenderer(
            link_meshes, object_meshes, args.width, args.height, bg_color=(0.0, 0.0, 0.0, 0.0),
        )

    renderer = _new_renderer()

    seg_id_map: dict[str, str] = {"0": "background_scene", "1": "robot"}
    for i, name in enumerate(object_names):
        seg_id_map[str(2 + i)] = name
    seg_pixel_counts = np.zeros(2 + len(object_names), dtype=np.int64)

    world_to_cam_per_frame: list[list[list[float]]] = []

    writer = VideoWriter(out_dir / "preview.mp4", fps=traj.fps)
    t_start = time.monotonic()
    try:
        for i in range(n_frames):
            if i > 0 and args.renderer_recycle > 0 and i % args.renderer_recycle == 0:
                renderer.close()
                renderer = _new_renderer()

            joints_i = traj.joint_positions[i] if traj.joint_positions is not None else None
            gripper_i = float(traj.gripper[i])
            link_poses = robot.link_poses(joints_i, gripper_i) if joints_i is not None else {}

            poses_i = (
                [traj.object_poses[i, j] for j in range(len(object_names))]
                if traj.object_poses is not None else []
            )

            if args.mode == "orbit":
                az = -args.orbit_degrees / 2.0 + args.orbit_degrees * i / max(n_frames - 1, 1)
                camera = free_camera(cloud_centre, orbit_distance, 180.0 + az, 18.0, K)
            else:
                camera = base_cam
            world_to_cam_per_frame.append(_world_to_cam_matrix(camera).tolist())

            combined, robot_mask, object_masks = renderer.render(link_poses, poses_i, camera)

            uv, cloud_depth = camera.project(cloud_xyz)
            canvas_bgr = np.full((args.height, args.width, 3), rps._BG, dtype=np.uint8)
            canvas_bgr = rps._draw_points(
                canvas_bgr, uv, cloud_depth, cloud_rgb, combined.depth, args.point_radius
            )
            canvas_bgr[combined.mask] = combined.color[:, :, ::-1][combined.mask]

            depth_m = _composite_depth(combined.depth, uv, cloud_depth, args.point_radius)
            depth_mm = np.clip(np.round(depth_m * 1000.0), 0, 65535).astype(np.uint16)

            seg = np.zeros((args.height, args.width), dtype=np.uint8)
            seg[robot_mask] = 1
            for j, m in enumerate(object_masks):
                seg[m] = 2 + j
            seg_pixel_counts[0] += int((seg == 0).sum())
            seg_pixel_counts[1] += int((seg == 1).sum())
            for j in range(len(object_names)):
                seg_pixel_counts[2 + j] += int((seg == 2 + j).sum())

            Image.fromarray(canvas_bgr[:, :, ::-1]).save(rgb_dir / f"{i:05d}.png")
            Image.fromarray(depth_mm).save(depth_dir / f"{i:05d}.png")
            Image.fromarray(seg).save(seg_dir / f"{i:05d}.png")
            writer.write(canvas_bgr)

            tcp_pose = _tcp_pose(link_poses) if link_poses else None
            frame_pose = {
                "frame": i,
                "timestamp_s": float(traj.timestamps[i]),
                "tcp_pose": tcp_pose.tolist() if tcp_pose is not None else None,
                "joint_positions": joints_i.tolist() if joints_i is not None else None,
                "gripper": gripper_i,
                "objects": {
                    name: (
                        traj.object_poses[i, j].tolist() if traj.object_poses is not None else None
                    )
                    for j, name in enumerate(object_names)
                },
            }
            (pose_dir / f"{i:05d}.json").write_text(json.dumps(frame_pose, indent=2))

            if i % 10 == 0 or i == n_frames - 1:
                elapsed = time.monotonic() - t_start
                logger.info(
                    "frame %d/%d  %.3fs/frame avg  RSS=%dMB",
                    i, n_frames, elapsed / (i + 1), _rss_mb(),
                )
    finally:
        writer.close()
        renderer.close()

    total_elapsed = time.monotonic() - t_start

    if args.action_json is not None and args.action_json.exists():
        action_payload: dict = json.loads(args.action_json.read_text())
        action_payload.setdefault("_source_file", str(args.action_json))
    else:
        action_payload = {
            "note": "no --action-json sidecar was provided (or the given path did not exist); "
            "this export has no recorded action spec.",
            "action_json_arg": str(args.action_json) if args.action_json else None,
        }
    (out_dir / "action.json").write_text(json.dumps(action_payload, indent=2, default=str))

    camera_json = {
        "fx": K.fx, "fy": K.fy, "cx": K.cx, "cy": K.cy, "width": K.width, "height": K.height,
        "world_to_cam_per_frame": world_to_cam_per_frame,
        "capture_episode": args.episode,
        "capture_camera_role": args.camera,
        "capture_camera_serial": serial,
        "capture_clip": args.clip,
        "capture_trajectory_row": capture_row,
        "mode": args.mode,
    }
    (out_dir / "camera.json").write_text(json.dumps(camera_json, indent=2))

    meta = {
        "frame_count": n_frames,
        "fps": traj.fps,
        "channels": {
            "rgb": "rgb/%05d.png, uint8 RGB (H, W, 3): robot+object mesh render composited "
            "over the carved background point cloud",
            "depth": "depth/%05d.png, uint16 (H, W) PNG, millimetres, 0 = no data. The "
            "renderer's mesh depth wherever the robot/object mesh is present (matching rgb "
            "and seg there), else the carved point-cloud depth, else 0.",
            "seg": "seg/%05d.png, uint8 (H, W) instance ids, see seg_id_map",
            "pose": "pose/%05d.json: tcp_pose (4x4 or null), joint_positions (7,) or null, "
            "gripper (float in [0, 1]), objects (name -> 4x4 world pose or null)",
        },
        "seg_id_map": seg_id_map,
        "seg_pixel_counts": {k: int(v) for k, v in zip(seg_id_map, seg_pixel_counts, strict=True)},
        "mode": args.mode,
        "renderer_recycle_every_n_frames": args.renderer_recycle,
        "runtime_seconds_total": total_elapsed,
        "runtime_seconds_per_frame_avg": total_elapsed / n_frames if n_frames else float("nan"),
        "input_warnings": traj.warnings,
        "object_assets": {
            asset.name: {
                "real_geometry": asset.is_real,
                "carved_from_background": asset.capture_transform is not None,
            }
            for asset in object_assets
        },
        "stand_in_input": bool(action_payload.get("STAND_IN", False)),
        "limitations": {
            "no_dynamics": "Kinematic replay only: no mass, friction, or inertia is modelled. A "
            "grasp is a rigid attachment of the object to the gripper frame, not a contact/force "
            "simulation. Commanded speed changes only the timing of the motion, never its path.",
            "single_viewpoint_capture": "The background is a single-viewpoint 2.5D point cloud "
            "unprojected from one real depth frame. Surfaces the capture camera never saw have no "
            "points; rendering far from the capture pose (e.g. --mode orbit) exposes real holes in "
            "the data, not a rendering bug.",
            "no_collision_checking": "No collision checking was performed between the robot, the "
            "manipulated object(s), or the static scene, at any stage of producing this trajectory "
            "or this render.",
            "no_settling_at_release": "A released object keeps exactly the pose it was holding at "
            "release; it does not fall, tip, or settle onto its support. Combined with the rigid "
            "grasp above, that means a released object can end up resting on an edge or floating "
            "clear of the surface it was placed on -- see placement_geometry below when the "
            "planner reports it. Levelling the object at release would be the one place a "
            "physics step is genuinely warranted (and is mass-independent), and is deliberately "
            "not done here.",
        },
    }
    # Surface the planner's own placement diagnostics rather than restating them:
    # if it measured a tilt it could not avoid, that number belongs next to the
    # limitation it illustrates, so a downstream consumer sees the defect's
    # magnitude and not just its existence.
    _place = action_payload.get("place_orientation_search")
    if isinstance(_place, dict) and "tilt_from_flat_deg" in _place:
        meta["placement_geometry"] = {
            "object_tilt_from_flat_deg": _place["tilt_from_flat_deg"],
            "why": "The placement target sits at the edge of the arm's reach and only converges "
            "with the wrist pitched off vertical; because the object is rigidly attached at its "
            "pick-time orientation, that pitch is transferred to the object and no reachable "
            "orientation lands it flat.",
            "orientation_search": _place,
        }
    (out_dir / "meta.json").write_text(json.dumps(meta, indent=2))

    print(f"wrote {n_frames} frames -> {out_dir.resolve()}")
    print(f"runtime: {total_elapsed:.1f}s total, {total_elapsed / n_frames:.3f}s/frame avg")
    print(f"peak RSS: {_rss_mb()} MB")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())