scripts/render_counterfactual_held.py

#!/usr/bin/env python
"""Counterfactual: after the grasp, the gripper never opens -- the brick rides in the
jaws for the rest of the episode instead of being placed in the drawer.

--- Why this is a different (and harder) experiment than the earlier gripper probe ---
A previous attempt (``scripts/render_counterfactual_grip.py``, a different episode)
overrode only the gripper *joint signal* and re-rendered the robot over the REAL
camera frame (``scripts/appearance_control.py``'s compositor). Measured there: the
resulting control differed from the real control by only 0.29-1.32% of pixels -- just
the two finger pads. The brick in that control was still the real, photographed brick
from the real video, so the generated clip showed the brick being released regardless
of what the rendered fingers did, and the experiment could not say whether the
generative model follows the control at all.

The fix is not a bigger override of the same kind -- it is rendering a DIFFERENT
BRICK TRAJECTORY. This script overrides the brick's *pose track*, not just the
gripper's, and composites everything (robot + repositioned brick + drawer) over
``master/plate.png`` -- a background plate with the real moving foreground already
removed (:mod:`fpgm.datagen.plate`) -- via the existing S8 exporter
(:class:`fpgm.datagen.export_vace.VaceExportStage`), which is exactly the "re-render
from poses + meshes, never from an observed photo" path S8 was built for (see that
module's own docstring). Nothing here draws over a real frame, so nothing here can
leak the real brick's real trajectory into the control.

--- What actually changes, precisely ---
1. ``master/events.json``'s ``brick.attached`` (S7, computed from poses alone --
   see :mod:`fpgm.datagen.events`) locates the grasp: the first frame the brick
   is rigidly attached to the gripper. This is a HEURISTIC symptom detector (that
   module's own words), but it is the one already used, in this exact codebase, to
   decide FK-attach bridging for occluded frames (``fpgm.datagen.object_poses.
   _bridge_gaps``) -- reusing it here keeps "grasp frame" consistent with what the
   rest of the pipeline already means by the term, rather than inventing a second
   detector.
2. From that frame on, ``T_gripper_brick = inv(T_world_gripper(grasp)) @
   T_world_brick(grasp)`` is fixed once, then ``T_world_brick(t) = T_world_gripper(t)
   @ T_gripper_brick`` for every later frame -- the exact rigid-attach formula
   ``_bridge_gaps`` already uses to bridge an occluded grasp span, applied here to
   EVERY frame from the grasp onward rather than only the occluded ones. Frames
   before the grasp keep the real, PnP-measured pose exactly.
3. ``T_world_gripper(t)`` comes from :meth:`fpgm.robot.kinematics.ArmKinematics.fk`
   on the REAL recorded arm-joint trajectory (untouched) and a gripper value that is
   real up to the grasp frame and FROZEN at its own closed value from there on --
   the same ``ArmKinematics.fk``-based TCP path S6/S7 already use to build
   ``T_world_tcp`` for ``attach_score`` (see ``fpgm.datagen.pipeline._run_s6``),
   so the FK feeding this override is the same FK that decided where the grasp is.
4. The brick's pose-source is ALSO overridden to ``FK_ATTACH`` (not left at whatever
   S6 originally recorded) from the grasp frame through the end of the episode.
   On the real footage, S6 could not solve the brick's pose past frame 87 (SAM3D
   PnP has nothing to solve once the brick is inside the drawer -- pose_source is
   GAP for 87-126, and S8 honestly does not render an object on a GAP frame). The
   counterfactual's whole premise is that the brick stays in view, gripped, so it
   MUST be rendered on those frames -- FK_ATTACH says so truthfully (this pose is a
   rigid-attach derivation, not a re-solved measurement), matching the enum's own
   documented meaning ("bridged via forward-kinematics while grasped").

--- What does NOT change, and the one place this leaves the scene inconsistent ---
The drawer is untouched: it keeps sliding on its own real, measured
(``drawer__T_world_obj``) track, including closing near the end of the episode, on
schedule, exactly as filmed. In the real episode the drawer closes *because* the
brick was just placed inside it and the arm is done with it; in this counterfactual
the brick never went in, yet the drawer still closes as if it had. This is a
deliberate scope limitation, not a bug: fixing it would mean inventing a
counterfactual drawer trajectory this episode's data says nothing about (the
episode's own robot never interacts with the drawer any differently once it is NOT
carrying a brick to it), which is a different, larger claim than "the gripper stays
closed." The rendered control is therefore physically self-consistent about the one
thing the intervention actually targets (gripper + brick) and openly inconsistent
about a second-order consequence (why would the drawer still close) it does not
attempt to model. Reported plainly here and in the run's own JSON report, not hidden.

--- Reuse, and why not more of it ---
Camera/robot/trajectory loading below mirrors ``fpgm.datagen.pipeline.run_s8_export``
line for line (same ``FlowsReader``/``EpisodeFrameIndex``, same extrinsics-candidate
selection, same ``RobotModel``/``ArmKinematics`` construction) because that IS the
production recipe for "S2's validated camera + S6's FK poses for this episode" and
reimplementing it differently here would risk silently disagreeing with the
``events.json`` this script depends on. :class:`~fpgm.datagen.export_vace.
VaceExportStage` itself is called unmodified -- the one thing overridden is the
in-memory ``object_T_world_obj["brick"]``/``object_pose_source["brick"]`` arrays
handed to it, exactly the extension point :meth:`VaceExportStage.run` already
exposes for a caller-supplied pose track. No second compositor is written.

**Why S8's per-window export, then a stitch, instead of a bespoke one-shot
renderer.** S8's window schedule (:func:`~fpgm.datagen.export_vace.compute_windows`)
is fixed at ``<=81`` frames -- Wan2.1-VACE's own ``4n+1, capped at 81`` training
contract -- so a single call cannot emit one continuous 127-frame control; this
episode always splits into windows ``[0,81) [40,121) [46,127)``. But
``scripts/sample_appearance_lora.py`` wants exactly one continuous
``(control.mp4, target.mp4)`` pair and does its OWN chunking at generation time
(``--chunk-frames``/``--chain``, confirmed to auto-chunk a >81-frame pair without the
caller pre-splitting). So this script exports the three S8 windows unmodified, then
stitches them back into one 127-frame pair by keeping, for each global frame, the
first window that ever covers it -- since a window's render depends only on that
single frame's own poses/camera/assets (no cross-frame filtering), any one of the
overlapping windows' renders of a shared frame is the same render.

**Which S8 channel becomes ``control.mp4``.** S8 writes five channels per window;
``control_vis.mp4`` -- "the renderer's own RGB, composited over the real plate" -- is
the one this project's own Wan-VACE ablation measured as the usable control (bg PSNR
22.99 vs 9-10 for the raw depth/seg/normal buffers, see ``export_vace.py``'s own
docstring), and it is also the only S8 channel that is an ordinary photograph-like
RGB video -- exactly ``AppearancePairDataset``'s ``control.mp4`` contract
(``fpgm.training.appearance_dataset``), unlike ``control_depth.mkv``/``control_seg.mkv``
(FFV1, integer-id semantics, a different dataset class entirely). Renaming
``control_vis.mp4`` to ``control.mp4`` in the stitched pair directory is the entire
translation needed between the two formats -- no new compositing.

--- Two honesty notes carried into every run's ``report.json`` ---
* ``VaceExportStage.run``'s on-disk cache fingerprint does not include object pose
  content (only uuid/camera/window range/caption/resolution/code_version) -- reusing
  the SAME output root for two different pose overrides would silently serve a stale
  render. This script always passes ``force=True`` and always uses a fresh,
  isolated ``out_root`` (never the episode's own ``outputs/datagen/<uuid>/<camera>/``
  tree, which the 9 production datagen episodes and 800 built training pairs
  already depend on).
* DiffSynth generation is NOT bit-reproducible across process launches even at a
  fixed seed (measured MAE ~2.0 between two launches of the same checkpoint on an
  identical control -- a floor this run's own report repeats verbatim so nobody
  reads a small step-500-vs-step-2000 difference as meaningful).

Usage (env ``fpgm`` -- this process does the pyrender/EGL render itself and only
shells out to env ``wan-train`` for the diffusion generation step, same
separation-of-concerns as ``render_counterfactual_grip.py``)::

    PYOPENGL_PLATFORM=egl PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python \\
        scripts/render_counterfactual_held.py --gpu 0
"""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import numpy as np

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

_DEFAULT_UUID = "AUTOLab+0d4edc83+2023-10-21-19h-07m-04s"
_DEFAULT_SERIAL = "22008760"
_WAN_PY = "/home/quang/miniconda3/envs/wan-train/bin/python"
_DEFAULT_CONFIG = REPO_ROOT / "configs" / "datagen_droid.yaml"

#: Pixel-value threshold (0-255, per channel, on decoded RGB) above which a pixel is
#: counted as "changed" between the factual and counterfactual control. Not zero,
#: because both controls independently pass through the same lossy H.264 encode/decode
#: (crf from ``DatagenConfig.control_video_crf``) even on frames with literally
#: identical geometry -- a bare ``!=`` would report a nonzero "difference" on every
#: single frame from compression noise alone. 10/255 is comfortably above typical H.264
#: banding at this crf and comfortably below "a brick rendered at a different 3D pose",
#: which moves a foreground region by tens to low-hundreds of pixel values, not single
#: digits.
_DIFF_PIXEL_THRESHOLD = 10


# --------------------------------------------------------------------------- #
# Episode context: mirrors fpgm.datagen.pipeline.run_s8_export's own loading
# --------------------------------------------------------------------------- #


def _load_context(uuid: str, camera_serial: str, config_path: Path, timer=None):
    """Load everything S8 already loads for this episode: camera, robot, trajectory.

    Deliberately the same recipe as ``fpgm.datagen.pipeline.run_s8_export`` (see this
    module's own docstring on why re-deriving a *different* camera/FK path here would
    risk silently disagreeing with ``events.json``, which was computed through this
    exact FK). Returns a plain dict rather than a new dataclass -- this script has one
    call site for each field, so a dataclass would only add indirection.
    """
    import h5py
    import numpy as np

    from fpgm.config_datagen import DatagenProfile
    from fpgm.data.droid_raw import read_mp4_properties
    from fpgm.data.pointworld import FlowsReader
    from fpgm.datagen.frame_index import EpisodeFrameIndex
    from fpgm.datagen.pipeline import _episode_gates_for_s8, _load_object_geometry
    from fpgm.datagen.robot_buffers import (
        TRAJECTORY_GRIPPER_POSITION_KEY,
        TRAJECTORY_JOINT_POSITIONS_KEY,
        load_extrinsics_candidates,
        read_native_camera_intrinsics,
    )
    from fpgm.robot.kinematics import ArmKinematics
    from fpgm.robot.urdf import RobotModel
    from fpgm.types import DataError

    with (timer.step("load_context") if timer else _noop()):
        profile = DatagenProfile.from_yaml(str(config_path))
        master_dir = profile.paths.master_dir(uuid, camera_serial)
        episode_dir = profile.paths.episode_dir(uuid)
        mp4_path = episode_dir / "recordings" / "MP4" / f"{camera_serial}.mp4"
        trajectory_path = episode_dir / "trajectory.h5"
        if not mp4_path.exists():
            raise DataError(f"mp4 not found: {mp4_path}")
        if not trajectory_path.exists():
            raise DataError(f"trajectory.h5 not found: {trajectory_path}")

        mp4_fps, n_video_frames, (video_w, video_h) = read_mp4_properties(mp4_path)
        with FlowsReader(profile.paths.flows_h5(uuid), episode_uuid=uuid) as reader:
            frame_index = EpisodeFrameIndex.build(
                uuid, reader, trajectory_path,
                mp4_properties=(mp4_fps, n_video_frames, (video_w, video_h)),
                camera_serial=camera_serial,
            )

        native_intrinsics = read_native_camera_intrinsics(
            profile.paths.flows_h5(uuid), camera_serial
        ).scaled(video_w, video_h)
        extrinsics_candidates = load_extrinsics_candidates(
            profile.paths.cameras_json(uuid), trajectory_path, camera_serial
        )
        s2_meta = json.loads((master_dir / "robot_buffers" / "meta.json").read_text())
        chosen_name = s2_meta["payload"]["chosen_extrinsics"]
        chosen = next(c for c in extrinsics_candidates if c.name == chosen_name)

        if not profile.paths.urdf.exists():
            raise DataError(f"URDF not found: {profile.paths.urdf}")
        robot = RobotModel(str(profile.paths.urdf), load_meshes=True)
        kin = ArmKinematics.from_robot_model(robot)
        with h5py.File(trajectory_path, "r") as f:
            joint_positions = np.asarray(f[TRAJECTORY_JOINT_POSITIONS_KEY])
            gripper_real = np.asarray(f[TRAJECTORY_GRIPPER_POSITION_KEY])

        import cv2

        with h5py.File(master_dir / "background_depth.h5", "r") as f:
            background_depth_mm = np.asarray(f["depth_mm"])
        plate_bgr = cv2.imread(str(master_dir / "plate.png"))
        if plate_bgr is None:
            raise DataError(f"could not read {master_dir / 'plate.png'}")
        background_plate_rgb = cv2.cvtColor(plate_bgr, cv2.COLOR_BGR2RGB)

        poses = np.load(master_dir / "poses.npz", allow_pickle=False)
        object_labels = tuple(str(x) for x in poses["labels"])
        object_T_world_obj = {label: poses[f"{label}__T_world_obj"] for label in object_labels}
        object_pose_source = {label: poses[f"{label}__pose_source"] for label in object_labels}
        object_meshes, point_cloud_objects = _load_object_geometry(
            master_dir / "meshes", list(object_labels)
        )

        episode_spec = json.loads((master_dir / "episode_spec.json").read_text())
        task_string = episode_spec.get("task") or f"robot manipulation episode {uuid}"

        events = json.loads((master_dir / "events.json").read_text())

        gates_passed = _episode_gates_for_s8(master_dir)

    return {
        "profile": profile, "uuid": uuid, "camera_serial": camera_serial,
        "master_dir": master_dir, "mp4_path": mp4_path, "mp4_fps": mp4_fps,
        "n_video_frames": n_video_frames, "video_w": video_w, "video_h": video_h,
        "frame_index": frame_index, "native_intrinsics": native_intrinsics,
        "world_to_cam": chosen.world_to_cam, "robot": robot, "kin": kin,
        "joint_positions": joint_positions, "gripper_real": gripper_real,
        "background_depth_mm": background_depth_mm, "background_plate_rgb": background_plate_rgb,
        "object_labels": object_labels, "object_T_world_obj": object_T_world_obj,
        "object_pose_source": object_pose_source, "object_meshes": object_meshes,
        "point_cloud_objects": point_cloud_objects, "task_string": task_string,
        "events": events, "gates_passed": gates_passed,
    }


class _noop:
    def __enter__(self):
        return None

    def __exit__(self, *a):
        return False


# --------------------------------------------------------------------------- #
# Grasp frame + rigid attachment
# --------------------------------------------------------------------------- #


def find_grasp_frame(events: dict, label: str = "brick") -> int:
    """First frame ``events.json``'s hysteresis-gated ``attached`` is True for ``label``.

    See the module docstring's point 1 for why ``events.json`` (S7, pose-derived,
    already computed) is reused rather than a second, script-local detector -- this
    is the exact same signal :func:`fpgm.datagen.object_poses._bridge_gaps` already
    conditions FK-attach bridging on.

    Raises:
        ValueError: If ``label`` is never attached in this episode -- nothing to
            hold onto for a "held" counterfactual.
    """
    attached = events[label]["attached"]
    for t, a in enumerate(attached):
        if a:
            return t
    raise ValueError(f"events.json: {label!r} is never attached in this episode")


def _resample_to_video_frames(
    row_indexed: np.ndarray, frame_index, n_video_frames: int
) -> np.ndarray:
    """``row_indexed`` (DROID trajectory-row axis) -> one value per video frame.

    Always through ``frame_index.trajectory_row`` -- never a raw fps-ratio division.
    ``fpgm.datagen.frame_index``'s own docstring documents this as a previously-real
    bug (rounding non-invertibility producing blank frames); every other loader in
    this pipeline (S6/S8's own ``joint_positions[row]``/``gripper[row]`` loops) goes
    through this exact same lookup, so this helper does too.
    """
    import numpy as np

    return np.asarray(
        [row_indexed[frame_index.trajectory_row(t)] for t in range(n_video_frames)],
        dtype=np.float64,
    )


def build_gripper_signal_and_world_gripper(
    ctx: dict, grasp_frame: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]:
    """Per-video-frame gripper values (real, and held-from-grasp) + T_world_gripper.

    ``T_world_gripper(t) = kin.fk(joint_positions[row(t)], gripper_cf(t))`` -- the
    same ``ArmKinematics.fk`` TCP path S6/S7 already use to build ``T_world_tcp`` for
    ``attach_score`` (see module docstring point 3). The arm-joint trajectory itself
    is the real recorded one at every frame, before and after the grasp -- only the
    gripper argument ever differs from the real signal.

    Returns:
        ``(gripper_cf, T_world_gripper, gripper_real_by_frame, closed_value)`` --
        the first three all length ``n_video_frames``; ``closed_value`` is the real
        signal's own value at the grasp frame (reported, not invented -- see this
        function's own measured note below). ``gripper_real_by_frame`` is returned
        too so the factual render can use the identical video-frame-indexed
        resampling as the counterfactual one, rather than a second, ad hoc mapping
        at the call site.
    """
    import numpy as np

    frame_index = ctx["frame_index"]
    n = ctx["n_video_frames"]
    joint_positions = ctx["joint_positions"]
    kin = ctx["kin"]

    gripper_real_by_frame = _resample_to_video_frames(ctx["gripper_real"], frame_index, n)
    closed_value = float(gripper_real_by_frame[grasp_frame])
    # MEASURED on this episode (not assumed): the recorded gripper signal is already
    # at its own plateau by the grasp frame -- 0.2599 -> 0.2863 between rows 43 and
    # 44, then flat at 0.2863 through row 60 -- so freezing at the grasp frame's own
    # value, rather than walking forward to a later plateau
    # (``render_counterfactual_grip.py``'s approach, needed there because ITS grasp
    # frame was picked from a bare threshold crossing mid-ramp), does not truncate
    # the real closing motion here: ``events.json``'s ``attached`` already only turns
    # True once ``attach_score``'s hysteresis (on_threshold=0.6) has cleared, which on
    # this episode happens to land past the ramp already.

    gripper_cf = gripper_real_by_frame.copy()
    gripper_cf[grasp_frame:] = closed_value

    t_world_gripper = np.empty((n, 4, 4), dtype=np.float64)
    for t in range(n):
        row = frame_index.trajectory_row(t)
        t_world_gripper[t] = kin.fk(joint_positions[row], float(gripper_cf[t]))
    return gripper_cf, t_world_gripper, gripper_real_by_frame, closed_value


def build_held_brick_track(
    ctx: dict, grasp_frame: int, t_world_gripper: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Rigidly attach the brick to the gripper from ``grasp_frame`` on.

    ``T_gripper_brick`` is fixed once, at the grasp frame's own real (PnP-measured)
    brick pose -- verified ``PoseSource.PNP`` below, not assumed, since the formula
    is meaningless applied to an already-fabricated pose. From there,
    ``T_world_brick(t) = T_world_gripper(t) @ T_gripper_brick`` for every later
    frame, exactly ``fpgm.datagen.object_poses._bridge_gaps``'s own FK-attach
    formula (see module docstring point 2), applied here through the END of the
    episode rather than only across one occluded gap.

    Returns:
        ``(t_world_brick_cf, pose_source_cf)``, both length ``n_video_frames``.
        Frames before ``grasp_frame`` are copied through unchanged, real values.
    """
    from fpgm.datagen.types import PoseSource
    from fpgm.geometry.transforms import invert_se3
    from fpgm.types import DataError

    t_world_brick_real = ctx["object_T_world_obj"]["brick"]
    pose_source_real = ctx["object_pose_source"]["brick"]
    if int(pose_source_real[grasp_frame]) != int(PoseSource.PNP):
        raise DataError(
            f"grasp frame {grasp_frame}: brick pose_source is "
            f"{PoseSource(int(pose_source_real[grasp_frame])).name}, not PNP -- "
            "T_gripper_brick would be anchored to a fabricated (not measured) pose"
        )

    t_gripper_brick = invert_se3(t_world_gripper[grasp_frame]) @ t_world_brick_real[grasp_frame]

    t_world_brick_cf = t_world_brick_real.copy()
    pose_source_cf = pose_source_real.copy()
    n = t_world_brick_real.shape[0]
    for t in range(grasp_frame, n):
        t_world_brick_cf[t] = t_world_gripper[t] @ t_gripper_brick
        pose_source_cf[t] = PoseSource.FK_ATTACH
    return t_world_brick_cf, pose_source_cf


# --------------------------------------------------------------------------- #
# S8 render (unmodified VaceExportStage) + stitch into one continuous pair
# --------------------------------------------------------------------------- #


def render_link_poses_for_frames(ctx: dict, gripper_by_frame: np.ndarray) -> list:
    """FK link poses for every video frame, given an already video-frame-indexed gripper trace.

    ``gripper_by_frame`` must already be resampled onto the video-frame axis (see
    :func:`_resample_to_video_frames`) -- this function only remaps ``joint_positions``
    (always trajectory-row-indexed) through ``frame_index.trajectory_row``, so the
    same call site works uniformly for the real signal (factual branch) and the
    frozen-from-grasp signal (counterfactual branch): the rendered fingers must match
    whichever ``T_world_gripper`` the brick was rigidly attached to, or the control
    would show the brick floating next to open jaws.
    """
    link_poses_by_frame = []
    for t in range(ctx["n_video_frames"]):
        row = ctx["frame_index"].trajectory_row(t)
        link_poses_by_frame.append(
            ctx["robot"].link_poses(ctx["joint_positions"][row], float(gripper_by_frame[t]))
        )
    return link_poses_by_frame


def render_windows_with_link_poses(
    ctx: dict, *, object_T_world_obj: dict, object_pose_source: dict,
    link_poses_by_frame: list, out_root: Path, caption: str, timer=None, step_name: str = "render",
) -> list:
    """Call the unmodified S8 exporter with a (possibly overridden) object pose track.

    The only things this function varies relative to ``fpgm.datagen.pipeline.
    run_s8_export`` are ``object_T_world_obj``/``object_pose_source``/
    ``link_poses_by_frame`` -- everything else (camera, robot meshes, background,
    gates, cfg) is this episode's own real, unmodified context. ``force=True``
    always: see module docstring's honesty note on why the on-disk cache
    fingerprint cannot tell a factual export apart from a counterfactual one at the
    same uuid/camera/window.
    """
    from fpgm.datagen.cache import StageCache
    from fpgm.datagen.export_vace import VaceExportStage

    with (timer.step(step_name, n=ctx["n_video_frames"]) if timer else _noop()):
        stage = VaceExportStage()
        samples = stage.run(
            uuid=ctx["uuid"], camera_serial=ctx["camera_serial"], mp4_path=ctx["mp4_path"],
            fps=ctx["mp4_fps"], n_video_frames=ctx["n_video_frames"],
            native_intrinsics=ctx["native_intrinsics"], world_to_cam=ctx["world_to_cam"],
            link_meshes=ctx["robot"].visual_meshes(), link_poses_by_frame=link_poses_by_frame,
            background_depth_mm=ctx["background_depth_mm"],
            background_plate_rgb=ctx["background_plate_rgb"], object_labels=ctx["object_labels"],
            object_meshes=ctx["object_meshes"], object_T_world_obj=object_T_world_obj,
            object_pose_source=object_pose_source, out_root=out_root, cfg=ctx["profile"].datagen,
            caption=caption, gates_passed=ctx["gates_passed"], cache=StageCache(out_root),
            force=True, caption_is_template_composed=True,
            point_cloud_objects=ctx["point_cloud_objects"],
        )
    return samples


def stitch_pair(
    samples: list, *, n_video_frames: int, uuid: str, camera_serial: str, caption: str,
    fps: float, video_wh: tuple[int, int], counterfactual_note: str, out_dir: Path,
) -> Path:
    """Stitch S8's overlapping per-window ``control_vis.mp4``/``target.mp4`` into one pair.

    See module docstring's "Reuse, and why not more of it" section for the full
    reasoning: for each global frame, the first window that covers it wins (any
    window's render of a shared frame is identical, since a frame's render depends
    only on that frame's own poses/camera/assets). Writes
    ``out_dir/{control.mp4,target.mp4,meta.json}`` -- the exact contract
    ``fpgm.training.appearance_dataset.AppearancePairDataset``/
    ``scripts/sample_appearance_lora.py --pair`` expects.
    """
    import cv2
    import numpy as np

    from fpgm.utils.io import ensure_dir

    ensure_dir(out_dir)
    windows = sorted(samples, key=lambda s: s.video_frame_start)
    w, h = video_wh
    control_frames: list = [None] * n_video_frames
    target_frames: list = [None] * n_video_frames
    covered = 0
    for s in windows:
        window_dir = s.target_video.parent
        cap_c = cv2.VideoCapture(str(window_dir / "control_vis.mp4"))
        cap_t = cv2.VideoCapture(str(window_dir / "target.mp4"))
        try:
            for i in range(s.video_frame_end - s.video_frame_start):
                ok_c, fc = cap_c.read()
                ok_t, ft = cap_t.read()
                if not (ok_c and ok_t):
                    raise RuntimeError(f"{window_dir}: short read at local frame {i}")
                g = s.video_frame_start + i
                if g >= covered:
                    control_frames[g] = fc
                    target_frames[g] = ft
        finally:
            cap_c.release()
            cap_t.release()
        covered = max(covered, s.video_frame_end)
    if covered < n_video_frames or any(f is None for f in control_frames):
        raise RuntimeError(
            f"stitch_pair: windows covered [0, {covered}) of {n_video_frames} -- "
            "compute_windows should always reach the episode end"
        )

    def _write(path: Path, frames: list) -> None:
        ff = "/home/quang/miniconda3/envs/ffmpeg_libs/bin/ffmpeg"
        p = subprocess.Popen(
            [ff, "-y", "-hide_banner", "-loglevel", "error", "-f", "rawvideo",
             "-pix_fmt", "bgr24", "-s", f"{w}x{h}", "-r", f"{fps:.4f}", "-i", "pipe:0",
             "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "14", str(path)],
            stdin=subprocess.PIPE,
        )
        for f in frames:
            p.stdin.write(np.ascontiguousarray(f).tobytes())
        p.stdin.close()
        if p.wait() != 0:
            raise RuntimeError(f"ffmpeg failed writing {path}")

    _write(out_dir / "control.mp4", control_frames)
    _write(out_dir / "target.mp4", target_frames)
    (out_dir / "meta.json").write_text(json.dumps({
        "uuid": uuid, "camera_serial": camera_serial, "caption": caption,
        "n_frames": n_video_frames, "fps": fps, "video_wh": [w, h],
        "counterfactual": counterfactual_note,
        "control_source": "fpgm.datagen.export_vace.VaceExportStage's control_vis "
                           "channel (robot+objects rendered, composited over master/plate.png), "
                           "stitched across S8's own overlapping window schedule -- see "
                           "scripts/render_counterfactual_held.py's module docstring.",
    }, indent=2))
    return out_dir / "meta.json"


# --------------------------------------------------------------------------- #
# Verify: how much does the counterfactual control actually differ, per frame
# --------------------------------------------------------------------------- #


def control_diff_fractions(factual_control: Path, counterfactual_control: Path) -> dict:
    """Per-frame fraction of pixels that differ (beyond H.264 noise) between two controls.

    This is the falsifiable check the module docstring exists to satisfy: the
    earlier gripper-only override measured 0.29-1.32% of pixels different (just the
    finger pads) because it never touched the brick. If this script's own numbers
    come out at that same ~1%, the control is NOT depicting the counterfactual and
    that must be reported as a failure, not glossed over on the way to a conclusion.
    """
    import cv2
    import numpy as np

    cap_a = cv2.VideoCapture(str(factual_control))
    cap_b = cv2.VideoCapture(str(counterfactual_control))
    fracs: list[float] = []
    try:
        while True:
            ok_a, fa = cap_a.read()
            ok_b, fb = cap_b.read()
            if not (ok_a and ok_b):
                break
            diff = np.abs(fa.astype(np.int16) - fb.astype(np.int16)).max(axis=2)
            fracs.append(float(np.mean(diff > _DIFF_PIXEL_THRESHOLD)))
    finally:
        cap_a.release()
        cap_b.release()
    if not fracs:
        raise RuntimeError("control_diff_fractions: no frames read from either control")
    arr = np.asarray(fracs)
    return {
        "pixel_threshold": _DIFF_PIXEL_THRESHOLD,
        "n_frames": len(fracs),
        "per_frame_diff_fraction": fracs,
        "mean_diff_fraction": float(arr.mean()),
        "max_diff_fraction": float(arr.max()),
        "min_diff_fraction": float(arr.min()),
    }


# --------------------------------------------------------------------------- #
# Generation (env wan-train, via subprocess) + side-by-side
# --------------------------------------------------------------------------- #


def _pick_gpu(min_free_gb: float = 10.0) -> int:
    """Best-effort: the GPU with the most free memory, refusing to guess below a floor.

    GPUs on this host are shared with other people (see the task brief) --
    ``nvidia-smi`` is queried at call time, never assumed from a config default, and
    a low-memory choice is reported loudly rather than silently attempted and OOMing
    a multi-minute generation job.
    """
    out = subprocess.run(
        ["nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader,nounits"],
        capture_output=True, text=True, check=True,
    ).stdout
    rows = [line.split(",") for line in out.strip().splitlines()]
    free_by_idx = {int(i): float(f) for i, f in rows}
    best = max(free_by_idx, key=free_by_idx.get)
    if free_by_idx[best] < min_free_gb * 1024:
        raise SystemExit(
            f"_pick_gpu: best candidate is GPU {best} with {free_by_idx[best]/1024:.1f} GB "
            f"free, below the {min_free_gb:.0f} GB floor -- pass --gpu explicitly to override"
        )
    return best


def generate(
    pair_dir: Path, out_dir: Path, *, gpu: int, num_frames: int, chunk_frames: int,
    steps: int, cfg_scale: float, seed: int, loras: list[Path], offload_text_encoder: bool,
    timer=None, step_name: str = "generate",
) -> None:
    """Two ``sample_appearance_lora.py`` launches: (baseline + loras[0]), then loras[1:].

    One process per checkpoint pairing rather than one process per checkpoint: the
    ``--baseline`` zero-shot arm only needs computing once per control (it does not
    depend on which LoRA step is being compared against it), and DiffSynth's model
    load is the dominant fixed cost per launch.
    """
    env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(gpu)}
    base_cmd = [
        _WAN_PY, str(REPO_ROOT / "scripts" / "sample_appearance_lora.py"),
        "--pair", str(pair_dir), "--out", str(out_dir),
        "--num-frames", str(num_frames), "--steps", str(steps),
        "--cfg-scale", str(cfg_scale), "--seed", str(seed),
        "--chunk-frames", str(chunk_frames), "--chain",
    ]
    if offload_text_encoder:
        base_cmd.append("--offload-text-encoder")

    with (timer.step(step_name, n=num_frames * len(loras)) if timer else _noop()):
        for i, lora in enumerate(loras):
            cmd = [*base_cmd, "--lora", str(lora)]
            if i == 0:
                cmd.append("--baseline")
            print("running:", " ".join(cmd), flush=True)
            result = subprocess.run(cmd, env=env, cwd=REPO_ROOT)
            if result.returncode != 0:
                raise SystemExit(
                    f"sample_appearance_lora.py failed (exit {result.returncode}): {lora}"
                )


def build_side_by_side(
    *, factual_control: Path, factual_output: Path, cf_control: Path, cf_output: Path,
    out_path: Path, fps: float, timer=None,
) -> None:
    """One row of four panels: factual control | factual output | cf control | cf output."""
    import cv2
    import numpy as np

    from fpgm.viz.video import H264Writer

    order = [
        ("factual control", factual_control), ("factual output", factual_output),
        ("counterfactual control", cf_control), ("counterfactual output", cf_output),
    ]
    caps = [cv2.VideoCapture(str(p)) for _, p in order]
    w = int(caps[0].get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(caps[0].get(cv2.CAP_PROP_FRAME_HEIGHT))
    n = min(int(c.get(cv2.CAP_PROP_FRAME_COUNT)) for c in caps)

    with (timer.step("side_by_side", n=n) if timer else _noop()):
        writer = H264Writer(out_path, w * 4, h, fps=round(fps), crf=18)
        try:
            for _ in range(n):
                tiles = []
                for (label, _path), cap in zip(order, caps, strict=True):
                    ok, frame_bgr = cap.read()
                    if not ok:
                        frame_bgr = np.zeros((h, w, 3), np.uint8)
                    tile = cv2.resize(frame_bgr, (w, h))[:, :, ::-1].copy()  # BGR -> RGB
                    cv2.putText(tile, label, (10, 26), cv2.FONT_HERSHEY_SIMPLEX,
                                0.65, (0, 0, 0), 4, cv2.LINE_AA)
                    cv2.putText(tile, label, (10, 26), cv2.FONT_HERSHEY_SIMPLEX,
                                0.65, (255, 255, 0), 1, cv2.LINE_AA)
                    tiles.append(tile)
                writer.write(np.hstack(tiles))
        finally:
            writer.close()
            for c in caps:
                c.release()


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--episode", default=_DEFAULT_UUID)
    ap.add_argument("--camera", default=_DEFAULT_SERIAL)
    ap.add_argument("--config", type=Path, default=_DEFAULT_CONFIG)
    ap.add_argument("--out-dir", type=Path, default=REPO_ROOT / "outputs" / "counterfactual_held")
    ap.add_argument("--gpu", type=int, default=None, help="default: auto-pick the freest GPU")
    ap.add_argument("--lora-step500", type=Path,
                    default=REPO_ROOT / "outputs" / "lora_appearance_2k" / "step-500.safetensors")
    ap.add_argument("--lora-step2000", type=Path,
                    default=REPO_ROOT / "outputs" / "lora_appearance_2k" / "step-2000.safetensors")
    ap.add_argument("--chunk-frames", type=int, default=81)
    ap.add_argument("--steps", type=int, default=30)
    ap.add_argument("--cfg-scale", type=float, default=5.0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--offload-text-encoder", action=argparse.BooleanOptionalAction, default=True)
    ap.add_argument("--skip-generate", action="store_true",
                    help="render + verify the controls and stop, for iterating on the "
                         "rigid-attach logic without paying for a wan-train generation")
    ap.add_argument("--skip-side-by-side", action="store_true")
    return ap.parse_args()


def main() -> int:
    args = parse_args()
    os.environ.setdefault("PYOPENGL_PLATFORM", "egl")

    from fpgm.utils.timing import StepTimer

    timer = StepTimer("counterfactual_held")

    render_gpu = args.gpu if args.gpu is not None else 0
    os.environ["EGL_DEVICE_ID"] = str(render_gpu)

    ctx = _load_context(args.episode, args.camera, args.config, timer=timer)
    n = ctx["n_video_frames"]

    grasp_frame = find_grasp_frame(ctx["events"], "brick")
    print(f"grasp frame (events.json brick.attached first True): {grasp_frame}", flush=True)

    with timer.step("build_attachment"):
        gripper_cf, t_world_gripper, gripper_real_by_frame, closed_value = (
            build_gripper_signal_and_world_gripper(ctx, grasp_frame)
        )
        t_world_brick_cf, pose_source_cf = build_held_brick_track(ctx, grasp_frame, t_world_gripper)
    print(
        f"gripper held closed from frame {grasp_frame} onward at value {closed_value:.4f} "
        f"(recorded signal's own value at the grasp frame; range was "
        f"[{float(ctx['gripper_real'].min()):.4f}, {float(ctx['gripper_real'].max()):.4f}] "
        "over the whole trajectory)", flush=True,
    )

    from fpgm.datagen.export_vace import compose_caption, describe_robot

    robot_description = describe_robot(
        base_link=ctx["robot"].base_link, link_names=ctx["robot"].link_names
    )
    caption = compose_caption(
        robot_description=robot_description, object_labels=ctx["object_labels"],
        task_string=ctx["task_string"],
    )

    out_dir = args.out_dir
    factual_link_poses = render_link_poses_for_frames(ctx, gripper_real_by_frame)
    cf_link_poses = render_link_poses_for_frames(ctx, gripper_cf)

    factual_samples = render_windows_with_link_poses(
        ctx, object_T_world_obj=ctx["object_T_world_obj"],
        object_pose_source=ctx["object_pose_source"], link_poses_by_frame=factual_link_poses,
        out_root=out_dir / "s8_render" / "factual", caption=caption, timer=timer,
        step_name="render_factual",
    )
    cf_object_T_world_obj = dict(ctx["object_T_world_obj"])
    cf_object_T_world_obj["brick"] = t_world_brick_cf
    cf_object_pose_source = dict(ctx["object_pose_source"])
    cf_object_pose_source["brick"] = pose_source_cf
    cf_samples = render_windows_with_link_poses(
        ctx, object_T_world_obj=cf_object_T_world_obj, object_pose_source=cf_object_pose_source,
        link_poses_by_frame=cf_link_poses, out_root=out_dir / "s8_render" / "counterfactual",
        caption=caption, timer=timer, step_name="render_counterfactual",
    )

    # The S8 export resolution (832x480, DatagenConfig.export_resolution), not the
    # native mp4 resolution -- read off the first window's own sample so the
    # stitcher never has to duplicate _check_resolution_contract's bucket logic.
    export_wh = tuple(factual_samples[0].resolution)

    with timer.step("stitch_pairs"):
        stitch_pair(
            factual_samples, n_video_frames=n, uuid=ctx["uuid"], camera_serial=ctx["camera_serial"],
            caption=caption, fps=ctx["mp4_fps"], video_wh=export_wh,
            counterfactual_note="none -- measured brick poses, real gripper signal (factual)",
            out_dir=out_dir / "pairs" / "factual",
        )
        stitch_pair(
            cf_samples, n_video_frames=n, uuid=ctx["uuid"], camera_serial=ctx["camera_serial"],
            caption=caption, fps=ctx["mp4_fps"], video_wh=export_wh,
            counterfactual_note=(
                f"gripper closes on the brick at frame {grasp_frame} and never opens again "
                f"(held at gripper value {closed_value:.4f}); brick rigidly attached to the "
                "gripper (T_world_brick(t) = T_world_gripper(t) @ T_gripper_brick) from that "
                "frame through the end of the episode. Drawer is UNCHANGED -- it still closes "
                "on its own real, measured schedule; see this script's module docstring."
            ),
            out_dir=out_dir / "pairs" / "counterfactual",
        )

    with timer.step("verify_control_diff"):
        diff = control_diff_fractions(
            out_dir / "pairs" / "factual" / "control.mp4",
            out_dir / "pairs" / "counterfactual" / "control.mp4",
        )
    print(
        f"control diff: mean {diff['mean_diff_fraction']*100:.2f}%  max "
        f"{diff['max_diff_fraction']*100:.2f}%  over {diff['n_frames']} frames "
        f"(threshold {diff['pixel_threshold']}/255 per channel)", flush=True,
    )
    if diff["mean_diff_fraction"] < 0.02:
        print(
            "WARNING: mean control diff is below 2% -- comparable to the earlier "
            "gripper-only-override failure mode (0.29-1.32%). This control may not be "
            "depicting the counterfactual; do not proceed to generation conclusions "
            "without checking the rendered frames directly.", flush=True,
        )

    report: dict = {
        "episode": ctx["uuid"], "camera_serial": ctx["camera_serial"],
        "grasp_frame": grasp_frame, "closed_gripper_value": closed_value,
        "caption": caption, "control_diff": diff,
        "export_resolution": list(export_wh),
    }

    if not args.skip_generate:
        gpu = args.gpu if args.gpu is not None else _pick_gpu()
        print(f"generation GPU: {gpu}", flush=True)
        loras = [args.lora_step500, args.lora_step2000]
        for branch in ("factual", "counterfactual"):
            generate(
                out_dir / "pairs" / branch, out_dir / "generation" / branch, gpu=gpu,
                num_frames=n, chunk_frames=args.chunk_frames, steps=args.steps,
                cfg_scale=args.cfg_scale, seed=args.seed, loras=loras,
                offload_text_encoder=args.offload_text_encoder, timer=timer,
                step_name=f"generate_{branch}",
            )

        if not args.skip_side_by_side:
            tag = f"{ctx['uuid']}__{ctx['camera_serial']}_f00000_chain{args.chunk_frames}"
            checkpoints = (("step-500", args.lora_step500), ("step-2000", args.lora_step2000))
            for step_name, lora in checkpoints:
                gen_factual = out_dir / "generation" / "factual" / f"{tag}__lora_{lora.stem}.mp4"
                gen_cf = out_dir / "generation" / "counterfactual" / f"{tag}__lora_{lora.stem}.mp4"
                side_by_side_path = out_dir / f"side_by_side_{step_name}.mp4"
                build_side_by_side(
                    factual_control=out_dir / "pairs" / "factual" / "control.mp4",
                    factual_output=gen_factual,
                    cf_control=out_dir / "pairs" / "counterfactual" / "control.mp4",
                    cf_output=gen_cf,
                    out_path=side_by_side_path, fps=ctx["mp4_fps"], timer=timer,
                )
                print(f"wrote {side_by_side_path}", flush=True)

    report["timings"] = timer.summary()
    (out_dir / "report.json").write_text(json.dumps(report, indent=2, default=str))
    print(timer.report(), flush=True)
    return 0


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