scripts/render_robot_random.py

#!/usr/bin/env python
"""Render the Franka+Robotiq arm waving through random joint poses, no input video.

    PYTHONPATH=src python scripts/render_robot_random.py --out outputs/robot_random.mp4

Unlike ``render_robot.py`` (which replays a real DROID trajectory over that
episode's camera footage), this samples a chaotic joint trajectory from
scratch -- a cubic spline through random waypoints per joint, each within that
joint's URDF ``<limit>`` range -- and renders it against a plain background
with a slowly orbiting camera. Written out as slow motion: the joint
trajectory is sampled at ``fps * slow_factor`` points across ``--duration``
seconds of "real" arm motion, then encoded at ``fps``, so the on-screen motion
plays back ``slow_factor``x slower than it actually happened.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
from scipy.interpolate import CubicSpline

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

from fpgm.geometry.camera import Camera  # noqa: E402
from fpgm.geometry.transforms import invert_se3  # noqa: E402
from fpgm.robot.render import RobotRenderer  # noqa: E402
from fpgm.robot.urdf import RobotModel  # noqa: E402
from fpgm.types import CameraIntrinsics  # noqa: E402
from fpgm.viz.overlays import VideoWriter  # noqa: E402

_DEFAULT_URDF = (
    REPO_ROOT
    / "third_party"
    / "robot_description"
    / "pointworld_franka_robotiq_2f85"
    / "franka_panda_robotiq_2f85.urdf"
)


def _random_joint_spline(
    rng: np.random.Generator, duration: float, waypoint_rate: float, arm_joint_limits: np.ndarray
):
    """Build a per-joint cubic spline through random waypoints within each joint's limits.

    A different random waypoint count/timing per joint (rather than one shared
    set of keyframe times) is what makes the motion look chaotic rather than a
    single rigid-body wave -- every joint independently changes direction at
    its own pace.
    """
    splines = []
    for lower, upper in arm_joint_limits:
        n_waypoints = max(2, int(rng.poisson(waypoint_rate * duration)) + 2)
        times = np.sort(rng.uniform(0.0, duration, size=n_waypoints))
        times[0], times[-1] = 0.0, duration
        values = rng.uniform(lower, upper, size=n_waypoints)
        splines.append(CubicSpline(times, values))
    return splines


def _random_gripper_spline(rng: np.random.Generator, duration: float, waypoint_rate: float):
    n_waypoints = max(2, int(rng.poisson(waypoint_rate * duration)) + 2)
    times = np.sort(rng.uniform(0.0, duration, size=n_waypoints))
    times[0], times[-1] = 0.0, duration
    values = rng.uniform(0.0, 1.0, size=n_waypoints)
    return CubicSpline(times, values)


def _look_at_world_to_cam(eye: np.ndarray, target: np.ndarray) -> np.ndarray:
    """Build an OpenCV-convention (+Z forward, +Y down) world-to-camera SE3 matrix."""
    world_up = np.array([0.0, 0.0, 1.0])
    forward = target - eye
    forward = forward / np.linalg.norm(forward)
    right = np.cross(forward, world_up)
    right = right / np.linalg.norm(right)
    down = np.cross(forward, right)

    cam_to_world = np.eye(4)
    cam_to_world[:3, 0] = right
    cam_to_world[:3, 1] = down
    cam_to_world[:3, 2] = forward
    cam_to_world[:3, 3] = eye
    return invert_se3(cam_to_world)


def _orbit_camera(t_frac: float, intrinsics: CameraIntrinsics, revolutions: float) -> Camera:
    """Camera slowly circling the robot's workspace, `t_frac` in [0, 1] over the clip."""
    angle = 2.0 * np.pi * revolutions * t_frac
    radius, elevation, target = 1.6, 0.55, np.array([0.0, 0.0, 0.5])
    eye = target + np.array([radius * np.cos(angle), radius * np.sin(angle), elevation])
    return Camera(intrinsics, _look_at_world_to_cam(eye, target))


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--urdf", type=Path, default=_DEFAULT_URDF)
    p.add_argument("--out", type=Path, default=REPO_ROOT / "outputs" / "robot_random.mp4")
    p.add_argument("--duration", type=float, default=5.0, help="seconds of simulated arm motion")
    p.add_argument("--fps", type=float, default=30.0, help="output video frame rate")
    p.add_argument("--slow-factor", type=float, default=5.0, help="slow-motion factor")
    p.add_argument("--width", type=int, default=480)
    p.add_argument("--height", type=int, default=480)
    p.add_argument("--waypoint-rate", type=float, default=1.0, help="random waypoints/second/joint")
    p.add_argument("--orbit-revolutions", type=float, default=0.5)
    p.add_argument("--seed", type=int, default=0)
    return p.parse_args()


def main() -> int:
    args = parse_args()
    if not args.urdf.exists():
        raise SystemExit(f"URDF not found: {args.urdf}")

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

    rng = np.random.default_rng(args.seed)
    joint_splines = _random_joint_spline(
        rng, args.duration, args.waypoint_rate, robot.arm_joint_limits()
    )
    gripper_spline = _random_gripper_spline(rng, args.duration, args.waypoint_rate)

    link_meshes = robot.visual_meshes()
    intrinsics = CameraIntrinsics(
        fx=0.9 * args.width, fy=0.9 * args.width,
        cx=args.width / 2.0, cy=args.height / 2.0,
        width=args.width, height=args.height,
    )

    # pyrender's OSMesa (software-GL) backend on this host leaks GL/Mesa state
    # every render() call regardless of scene reuse (~180MB/frame at 480x480,
    # confirmed with a standalone probe) -- recreating the renderer every
    # `_RENDERER_RECYCLE` frames bounds that growth instead of letting a long
    # clip OOM the (shared) host.
    _RENDERER_RECYCLE = 10

    def _new_renderer():
        return RobotRenderer(link_meshes, args.width, args.height, bg_color=(0.05, 0.05, 0.08, 1.0))

    renderer = _new_renderer()

    n_frames = int(round(args.duration * args.fps * args.slow_factor))
    sim_times = np.linspace(0.0, args.duration, n_frames, endpoint=False)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    writer = VideoWriter(args.out, fps=args.fps)
    try:
        for i, t in enumerate(sim_times):
            if i > 0 and i % _RENDERER_RECYCLE == 0:
                renderer.close()
                renderer = _new_renderer()
            joints = np.array([s(t) for s in joint_splines])
            gripper = float(np.clip(gripper_spline(t), 0.0, 1.0))
            link_poses = robot.link_poses(joints, gripper)
            camera = _orbit_camera(i / n_frames, intrinsics, args.orbit_revolutions)
            result = renderer.render(link_poses, camera)
            writer.write(result.color[:, :, ::-1])  # RGB -> BGR
            if i % 50 == 0:
                print(f"frame {i}/{n_frames}  t={t:.2f}s  mask coverage={result.mask.mean():.3f}")
    finally:
        writer.close()
        renderer.close()

    playback_seconds = n_frames / args.fps
    print("\n" + "=" * 78)
    print(f"frames written:   {n_frames}")
    print(f"simulated motion: {args.duration:.1f}s -> playback: {playback_seconds:.1f}s "
          f"({args.slow_factor:.0f}x slow motion)")
    print(f"output: {args.out}")
    print("=" * 78)
    return 0


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