scripts/render_counterfactual_grip.py

#!/usr/bin/env python
"""Counterfactual: the gripper closes on the brick and never opens again.

Does the appearance LoRA actually follow the rendered control, or has it
memorised this one episode? Episode ``AUTOLab+0d4edc83+2023-10-21-19h-37m-47s``,
camera ``22008760`` ("put brick in drawer shelf and close drawer"), is the LoRA's
own holdout pair (``outputs/appearance_holdout/.../``): the real footage shows
the gripper closing on a brick partway through, carrying it over the drawer, and
opening to drop it in. If we render a control where the gripper closes on the
brick and **never opens again**, and generate from it with the LoRA, one of two
things happens:

  * the generated brick stays gripped, following the control -- the LoRA is
    doing what it is supposed to (appearance transfer, not clip reproduction);
  * the generated brick gets released anyway, on cue with the real footage --
    the LoRA has memorised *this episode's* release event and is replaying it
    regardless of what the control says.

Either result is reportable; neither is assumed going in.

--- Honesty requirement #1: the composite is internally inconsistent by design ---
``appearance_control.py``'s control is a URDF-rendered Franka composited over
the **real** camera frame (see its module docstring) -- there is no clean
background plate for this pair, only the real video with the real robot painted
over. Overriding the gripper changes the *rendered* robot only. The background
underneath is still the real frame: the real brick still visibly leaves the
gripper and drops into the drawer, the real drawer still visibly closes, on
schedule, regardless of what the rendered robot's fingers are doing. So
``control_counterfactual.mp4`` is not a physically consistent scene -- it is a
robot that (by construction) holds while the photograph behind it shows the
object being let go. This is a limitation of testing a gripper counterfactual on
an appearance pair, not a hidden bug, and it directly shapes how the generated
video has to be read: if the LoRA follows the real background rather than the
rendered robot, that is not surprising on physical-plausibility grounds, but it
is exactly the memorisation failure mode this test is built to expose.

--- Honesty requirement #2: the real arm/gripper can show through -----------
The composite only draws the rendered robot inside its own silhouette; every
other pixel is the untouched real frame, real robot included. Once the gripper
override activates, the counterfactual render's finger silhouette stops
tracking the real recorded gripper's silhouette (the real hand keeps opening
and moving through the second grasp cycle and retreat; the rendered one sits
frozen half-closed). Wherever the real gripper's true footprint extends beyond
the counterfactual render's, the real fingers leak through around the rendered
ones. This script quantifies that directly: it renders the real recorded
gripper trace as a second, throwaway pass alongside the overridden one (see
``ClipBuilder.render_loop``'s ``gripper_reference`` argument, added to
``appearance_control.py`` for exactly this) and measures the silhouette
difference per frame. See ``gripper_override.json``'s ``arm_showthrough`` block
and the run's printed report for the actual numbers -- they vary by episode and
threshold and are not asserted here.

--- The override itself: measured, not assumed ------------------------------
The recorded ``gripper_position`` signal is **not** normalised the same way
across episodes in absolute terms -- ``RobotModel`` treats it as [0, 1] with 0
open, but a demo's peak commanded closure depends on the grasped object's width
(the arm stops closing on contact, not at a fixed target). A sibling AUTOLab
episode's closure tops out at 0.286; THIS episode's first grasp plateaus at
~0.64 and a later, unrelated closure (after the brick is already released,
almost certainly the arm gripping empty air near a drawer-closing motion)
reaches ~0.87. Reusing an absolute threshold from one episode on another would
therefore be wrong, and the module docstring's own honesty bar requires
verifying this episode's own values rather than assuming them -- so the
threshold below is defined relative to *this episode's own* observed range and
every number it produces is computed at runtime and logged, never hard-coded.

Algorithm (fully generic -- no episode-specific branch, works on any episode):

  1. ``baseline = gripper[0]``, ``ep_max = gripper.max()`` (this episode's own
     open/closed extremes).
  2. ``threshold = baseline + close_frac * (ep_max - baseline)``, default
     ``close_frac=0.5``. Chosen because it is simple, round, and -- checked
     against 0.3 and 0.7 on this episode -- lands on the exact same fully-closed
     plateau either way (frame 103, value 0.6388); the threshold only has to be
     high enough to reject the ramp's shallow opening seconds, not pin an exact
     row, so 50% has margin on both sides without being fussy about it.
  3. ``frame_threshold_crossed`` = first row where ``gripper >= threshold``.
     This *locates* the closing event -- it is not where the override starts.
  4. From there, walk forward while the signal is non-decreasing (allowing for
     float noise) to the first local maximum/plateau: ``frame_hold_start``.
     This is where the real closing motion actually finishes.
  5. ``closed_value = gripper[frame_hold_start]``; every row from
     ``frame_hold_start`` onward is overwritten with this constant. Rows before
     it are untouched, so the real, recorded closing animation plays out
     exactly as filmed and only *freezes* once the grasp is complete, rather
     than snapping early mid-ramp (which step 3's frame alone would do -- tried
     first, rejected: it visibly jerks the render straight to fully-closed a
     few frames before the real motion would have gotten there, for no benefit
     over just letting the real ramp finish and holding after).
  6. That constant is held through the entire rest of the episode, silently
     overriding the real signal's later partial release, second regrasp, and
     final open -- deliberately: "never opens again" means the counterfactual
     ignores everything the real trajectory's gripper channel does after this
     point, not just the release event the episode is named for.

Only the gripper argument to ``RobotModel.link_poses`` changes. The arm joint
trajectory (``joint_positions``) is read from ``trajectory.h5`` and passed
through untouched -- the arm still reaches, still moves toward the drawer,
exactly as recorded.

--- Pipeline reuse -----------------------------------------------------------
Depth is stage 1 of the existing pipeline (``appearance_depth.py``, env
``moge``) and must already exist in ``--work`` before this runs -- this script
does not span conda environments, it only renders (env ``fpgm``/EGL) and then
shells out to ``sample_appearance_lora.py`` (env ``wan-train``) for generation,
same separation-of-concerns the rest of this pipeline already uses.
Rendering itself is ``appearance_control.ClipBuilder.load_clip`` +
``.render_loop``, unmodified except for the added optional ``gripper_reference``
diff pass -- no rendering/compositing logic is duplicated here.

Usage (three stages, three envs, matching the rest of this pipeline)::

    # stage 1 (moge), only if outputs/counterfactual_grip/work is missing the
    # depth artifacts for this episode -- regenerate them if deleted:
    CUDA_VISIBLE_DEVICES=<gpu> /home/quang/miniconda3/envs/moge/bin/python -u \\
        scripts/appearance_depth.py --work outputs/counterfactual_grip/work \\
        --uuid AUTOLab+0d4edc83+2023-10-21-19h-37m-47s

    # stage 2+3 (fpgm render, then wan-train generate, orchestrated here):
    PYOPENGL_PLATFORM=egl PYTHONPATH=src /home/quang/miniconda3/envs/fpgm/bin/python \\
        scripts/render_counterfactual_grip.py --gpu <gpu>
"""

from __future__ import annotations

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

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

_DEFAULT_UUID = "AUTOLab+0d4edc83+2023-10-21-19h-37m-47s"
_DEFAULT_SERIAL = "22008760"
_WAN_PY = "/home/quang/miniconda3/envs/wan-train/bin/python"


def gripper_override(gripper, close_frac: float) -> dict:
    """Locate the first real closing event and pin the signal there forever.

    See the module docstring's "The override itself" section for the full
    rationale. Returns every intermediate number (not just the final array) so
    the caller can log and persist them -- this is the one place in the script
    where "document the chosen threshold and why" has to be backed by numbers
    computed from THIS episode's signal, not asserted.
    """
    import numpy as np

    g = np.asarray(gripper, dtype=np.float64)
    baseline = float(g[0])
    ep_max = float(g.max())
    if ep_max <= baseline + 1e-6:
        raise ValueError(
            f"gripper signal never closes measurably above its own baseline "
            f"({baseline:.4f}) -- nothing to grab onto for an override"
        )
    threshold = baseline + close_frac * (ep_max - baseline)
    crossed = np.flatnonzero(g >= threshold)
    if crossed.size == 0:
        raise ValueError(f"gripper never reaches the {close_frac:.0%} threshold {threshold:.4f}")
    t_cross = int(crossed[0])
    j = t_cross
    while j + 1 < len(g) and g[j + 1] >= g[j] - 1e-9:
        j += 1
    t_hold = j
    closed_value = float(g[t_hold])
    overridden = g.copy()
    overridden[t_hold:] = closed_value
    return {
        "baseline": baseline,
        "episode_max": ep_max,
        "close_frac": close_frac,
        "threshold": threshold,
        "frame_threshold_crossed": t_cross,
        "value_at_threshold_crossed": float(g[t_cross]),
        "frame_hold_start": t_hold,
        "closed_value": closed_value,
        "n_rows_total": int(len(g)),
        "n_rows_overridden": int(len(g) - t_hold),
        "gripper": overridden,
    }


def summarize_showthrough(diff: dict, frame_wh: tuple[int, int], px_floor: int = 50) -> dict:
    """Aggregate the per-frame real-vs-counterfactual silhouette diff.

    ``diff["extra_real_px"][t]`` is pixels where the REAL recorded gripper's
    render would have drawn something that the COUNTERFACTUAL render's own
    silhouette does not cover this frame -- i.e. the composite would leave the
    untouched real frame (real arm included) visible there instead of painting
    over it. This is an estimate, not a segmentation of the photograph: it uses
    the real-gripper render's own silhouette as a proxy for "where the real arm
    is on screen", which is exactly the assumption the rest of this pipeline
    already relies on (the MoGe depth-alignment fit in ``render_and_composite``
    only makes sense if the render tracks the real arm's screen position).
    """
    import numpy as np

    extra = np.asarray(diff["extra_real_px"], dtype=np.int64)
    drawn_real = np.asarray(diff["drawn_real_px"], dtype=np.int64)
    frame_area = frame_wh[0] * frame_wh[1]
    affected = extra > px_floor
    return {
        "n_frames": int(len(extra)),
        "px_floor": px_floor,
        "n_frames_with_showthrough": int(affected.sum()),
        "frac_frames_with_showthrough": round(float(affected.mean()), 4),
        "mean_extra_px_all_frames": round(float(extra.mean()), 1),
        "mean_extra_px_affected_frames": (
            round(float(extra[affected].mean()), 1) if affected.any() else 0.0
        ),
        "mean_extra_frac_of_frame_on_affected_frames": (
            round(float(extra[affected].mean() / frame_area), 5) if affected.any() else 0.0
        ),
        "total_extra_px_over_total_real_drawn_px": round(
            float(extra.sum() / max(int(drawn_real.sum()), 1)), 5
        ),
    }


def write_grid(panels: dict, out: Path, fps: float, timer=None) -> None:
    """2x2 side-by-side: real control/output on top, counterfactual on bottom.

    Reads all four videos with OpenCV (they are already frame-aligned: same
    pair, same ``--num-frames``/``--start``/seed on the generation side), labels
    each panel, and pipes the combined frame to the same ffmpeg-subprocess
    writer ``appearance_control.H264Writer`` uses -- no new video-writing code.
    """
    import cv2
    import numpy as np
    from appearance_control import H264Writer

    order = ["real_control", "real_lora", "cf_control", "cf_lora"]
    labels = {
        "real_control": "REAL control",
        "real_lora": "REAL + LoRA step-500",
        "cf_control": "COUNTERFACTUAL control (grip held)",
        "cf_lora": "COUNTERFACTUAL + LoRA step-500",
    }
    caps = {k: cv2.VideoCapture(str(panels[k])) for k in order}
    w = int(caps["real_control"].get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(caps["real_control"].get(cv2.CAP_PROP_FRAME_HEIGHT))
    n = min(int(c.get(cv2.CAP_PROP_FRAME_COUNT)) for c in caps.values())

    step_ctx = timer.step("side_by_side", n=n) if timer else _noop()
    with step_ctx:
        writer = H264Writer(out, w * 2, h * 2, fps, crf=18)
        try:
            for _ in range(n):
                tiles = []
                for k in order:
                    ok, frame = caps[k].read()
                    if not ok:
                        frame = np.zeros((h, w, 3), np.uint8)
                    tile = frame.copy()
                    cv2.putText(tile, labels[k], (10, 26), cv2.FONT_HERSHEY_SIMPLEX,
                               0.7, (0, 0, 0), 4, cv2.LINE_AA)
                    cv2.putText(tile, labels[k], (10, 26), cv2.FONT_HERSHEY_SIMPLEX,
                               0.7, (0, 255, 255), 1, cv2.LINE_AA)
                    tiles.append(tile)
                top = np.hstack([tiles[0], tiles[1]])
                bottom = np.hstack([tiles[2], tiles[3]])
                writer.write(np.vstack([top, bottom]))
        finally:
            writer.close()
            for c in caps.values():
                c.release()


class _noop:
    def __enter__(self):
        return None

    def __exit__(self, *a):
        return False


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("--work", type=Path, default=REPO_ROOT / "outputs/counterfactual_grip/work",
                    help="MoGe depth artifacts (stage 1); must already exist")
    ap.add_argument("--out", type=Path, default=REPO_ROOT / "outputs/counterfactual_grip")
    ap.add_argument("--episodes-root", type=Path, default=REPO_ROOT / "data/droid_raw")
    ap.add_argument("--cameras-dir", type=Path,
                    default=REPO_ROOT / "data/pointworld/droid/cameras")
    ap.add_argument("--intrinsics", type=Path,
                    default=REPO_ROOT / "configs/droid_camera_intrinsics.json")
    ap.add_argument("--urdf", type=Path, default=None)
    ap.add_argument("--depth-tol-m", type=float, default=0.15)
    ap.add_argument("--crf", type=int, default=18)
    ap.add_argument("--fps", type=float, default=15.0)
    ap.add_argument("--close-frac", type=float, default=0.5,
                    help="threshold as a fraction of this episode's own "
                         "[gripper[0], gripper.max()] range; see module docstring")
    ap.add_argument("--gpu", type=int, default=2,
                    help="sets EGL_DEVICE_ID for rendering and CUDA_VISIBLE_DEVICES "
                         "for the wan-train generation subprocess")
    ap.add_argument("--lora", type=Path,
                    default=REPO_ROOT / "outputs/lora_appearance_2k/step-500.safetensors")
    ap.add_argument("--num-frames", type=int, default=249,
                    help="matches the existing real-pair generation "
                         "(outputs/lora_samples_rigid) for a frame-aligned comparison")
    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("--chunk-frames", type=int, default=81)
    ap.add_argument("--real-samples-dir", type=Path,
                    default=REPO_ROOT / "outputs/lora_samples_rigid",
                    help="existing REAL control/LoRA-output bundle to reuse for "
                         "the side-by-side, instead of re-measuring a baseline "
                         "that is already on disk")
    ap.add_argument("--skip-generate", action="store_true",
                    help="render control_counterfactual.mp4 and stop -- for "
                         "iterating on the gripper-override logic without "
                         "paying for a wan-train generation each time")
    ap.add_argument("--skip-grid", action="store_true")
    return ap.parse_args()


def main() -> int:
    args = parse_args()
    os.environ["EGL_DEVICE_ID"] = str(args.gpu)
    os.environ.setdefault("PYOPENGL_PLATFORM", "egl")

    from fpgm.utils.timing import StepTimer

    timer = StepTimer("counterfactual_grip")

    import appearance_control as ac
    import numpy as np

    args.out.mkdir(parents=True, exist_ok=True)
    depth_json = args.work / f"{args.episode}__{args.camera}.depth.json"
    depth_npy = args.work / f"{args.episode}__{args.camera}.depth.npy"
    if not (depth_json.exists() and depth_npy.exists()):
        raise SystemExit(
            f"missing MoGe depth for {args.episode}/{args.camera} in {args.work} -- "
            f"regenerate with scripts/appearance_depth.py (env moge) first; see "
            f"this script's module docstring for the exact command"
        )
    depth_meta = json.loads(depth_json.read_text())

    urdf = args.urdf
    if urdf is None:
        from fpgm.config_datagen import DatagenProfile
        urdf = Path(DatagenProfile.from_yaml(
            str(REPO_ROOT / "configs/datagen_droid.yaml")).paths.urdf)

    builder = ac.ClipBuilder(
        urdf, ac.load_intrinsic_table(args.intrinsics), args.depth_tol_m, args.crf)
    try:
        with timer.step("load_clip"):
            clip = builder.load_clip(depth_meta, args.episodes_root, args.cameras_dir)
            real_gripper = clip["gripper"]

        with timer.step("gripper_override"):
            override = gripper_override(real_gripper, args.close_frac)
        print(
            f"gripper signal (episode {args.episode}, n_rows={override['n_rows_total']}): "
            f"baseline={override['baseline']:.4f} episode_max={override['episode_max']:.4f} "
            f"threshold({args.close_frac:.0%})={override['threshold']:.4f}\n"
            f"  first crosses threshold at row {override['frame_threshold_crossed']} "
            f"(value {override['value_at_threshold_crossed']:.4f})\n"
            f"  natural close completes (hold starts) at row {override['frame_hold_start']} "
            f"(closed_value={override['closed_value']:.4f}), "
            f"{override['n_rows_overridden']} of {override['n_rows_total']} rows overridden",
            flush=True,
        )

        pair_dir = args.out / "pair"
        pair_dir.mkdir(parents=True, exist_ok=True)
        depth_all = np.load(args.work / f"{args.episode}__{args.camera}.depth.npy",
                            mmap_mode="r")
        with timer.step("render_counterfactual", n=clip["n"]):
            stats = builder.render_loop(
                clip["joint_positions"], override["gripper"], clip["mp4"], clip["camera"],
                clip["frames_per_step"], clip["n"], depth_all, clip["out_w"], clip["out_h"],
                pair_dir, args.fps, gripper_reference=real_gripper,
            )

        ep_meta = json.loads((args.episodes_root / args.episode / "metadata.json").read_text())
        pair_meta = {
            "uuid": args.episode, "camera_serial": args.camera, "n_frames": len(stats["a"]),
            "fps": args.fps, "video_wh": [clip["out_w"], clip["out_h"]],
            "caption": ep_meta.get("current_task"), "lab": ep_meta.get("lab"),
            "control": "control.mp4", "target": "target.mp4",
            "counterfactual": "gripper closes on the brick and never opens again "
                              "(see gripper_override.json / module docstring)",
            "depth_tol_m": args.depth_tol_m,
            "occluded_px_fraction": round(
                float(np.sum(stats["reject"]) / max(int(np.sum(stats["robot_px"])), 1)), 5),
        }
        (pair_dir / "meta.json").write_text(json.dumps(pair_meta, indent=2))

        control_cf = args.out / "control_counterfactual.mp4"
        control_cf.write_bytes((pair_dir / "control.mp4").read_bytes())

        showthrough = summarize_showthrough(
            stats["diff"], (clip["out_w"], clip["out_h"]))
        print(
            f"real-arm show-through: {showthrough['n_frames_with_showthrough']}/"
            f"{showthrough['n_frames']} frames affected "
            f"({showthrough['frac_frames_with_showthrough']*100:.1f}%), "
            f"mean {showthrough['mean_extra_frac_of_frame_on_affected_frames']*100:.2f}% "
            f"of frame area on affected frames", flush=True,
        )

        override_out = dict(override)
        override_out.pop("gripper")
        override_out["arm_showthrough"] = showthrough
        (args.out / "gripper_override.json").write_text(json.dumps(override_out, indent=2))
    finally:
        builder.close()

    if args.skip_generate:
        print(timer.report(), flush=True)
        return 0

    tag = f"{args.episode}__{args.camera}_f00000_chain{args.chunk_frames}"
    gen_env = {**os.environ, "CUDA_VISIBLE_DEVICES": str(args.gpu)}
    gen_cmd = [
        _WAN_PY, str(REPO_ROOT / "scripts/sample_appearance_lora.py"),
        "--pair", str(pair_dir), "--out", str(args.out),
        "--lora", str(args.lora), "--baseline",
        "--num-frames", str(args.num_frames), "--steps", str(args.steps),
        "--cfg-scale", str(args.cfg_scale), "--seed", str(args.seed),
        "--chunk-frames", str(args.chunk_frames), "--chain",
    ]
    with timer.step("generate_wan", n=2 * args.num_frames):
        print("running:", " ".join(gen_cmd), flush=True)
        result = subprocess.run(gen_cmd, env=gen_env, cwd=REPO_ROOT)
    if result.returncode != 0:
        print(timer.report(), flush=True)
        raise SystemExit(f"sample_appearance_lora.py failed (exit {result.returncode})")

    cf_lora = args.out / f"{tag}__lora_{args.lora.stem}.mp4"
    cf_baseline = args.out / f"{tag}__baseline.mp4"
    cf_control_249 = args.out / f"{tag}__control.mp4"
    print(f"wrote {cf_lora}\nwrote {cf_baseline}", flush=True)

    if not args.skip_grid:
        real_tag = f"{args.episode}__{args.camera}_f00000_chain{args.chunk_frames}"
        real_control = args.real_samples_dir / f"{real_tag}__control.mp4"
        real_lora = args.real_samples_dir / f"{real_tag}__lora_{args.lora.stem}.mp4"
        if real_control.exists() and real_lora.exists() and cf_lora.exists():
            write_grid(
                {"real_control": real_control, "real_lora": real_lora,
                 "cf_control": cf_control_249, "cf_lora": cf_lora},
                args.out / "side_by_side.mp4", args.fps, timer=timer,
            )
            print(f"wrote {args.out / 'side_by_side.mp4'}", flush=True)
        else:
            print(
                f"skipping side-by-side: missing one of {real_control}, {real_lora}, "
                f"{cf_lora} -- pass --real-samples-dir or rerun without --skip-generate",
                flush=True,
            )

    print(timer.report(), flush=True)
    return 0


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