lerobot/benchmarks/video/benchmark_video_drift.py

#!/usr/bin/env python
"""Benchmark the timestamp drift produced by the *actual* codebase recording path.

Unlike the simulation in ``tests/datasets/test_video_drift.py``
(``test_round6_accumulates_drift_but_actual_duration_does_not``), this script does not
re-implement any arithmetic. It records episodes through the real
``LeRobotDataset.create / add_frame / save_episode / finalize`` pipeline (PNG -> mp4
encoding + ``concatenate_video_files``), then measures how far the ``from_timestamp`` /
``to_timestamp`` values stored in the episode metadata drift from the PTS actually
decoded from the concatenated video file.

Drift sources exercised here:
- float accumulation of ``to_timestamp = from_timestamp + ep_duration``
- per-episode ``get_video_duration_in_s`` vs the frame's real PTS after concatenation

Run:
    python benchmarks/video/benchmark_video_drift.py
    python benchmarks/video/benchmark_video_drift.py --fps 30 --num-episodes 500
"""

import argparse
import shutil
import tempfile
from pathlib import Path

import av
import numpy as np

from lerobot.datasets.io_utils import load_episodes
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.video_utils import get_video_duration_in_s

VIDEO_KEY = "observation.images.laptop"


def _decode_all_frame_pts(video_path: Path | str) -> list[float]:
    """Return the PTS (seconds) of every frame in decode order, in a single pass."""
    with av.open(str(video_path)) as container:
        stream = container.streams.video[0]
        time_base = stream.time_base
        return [float(frame.pts * time_base) for frame in container.decode(stream)]


def _record_dataset(
    root: Path,
    fps: int,
    frames_per_episode: list[int],
    streaming: bool,
) -> LeRobotDataset:
    features = {
        VIDEO_KEY: {"dtype": "video", "shape": (3, 64, 96), "names": ["channels", "height", "width"]},
        "state": {"dtype": "float32", "shape": (2,), "names": None},
    }
    dataset = LeRobotDataset.create(
        repo_id="benchmark/video_drift",
        fps=fps,
        features=features,
        root=root,
        streaming_encoding=streaming,
        # Force every episode into a single concatenated video file.
        video_files_size_in_mb=10_000,
    )
    rng = np.random.RandomState(0)
    for n_frames in frames_per_episode:
        for _ in range(n_frames):
            dataset.add_frame(
                {
                    VIDEO_KEY: rng.randint(0, 256, (64, 96, 3), dtype=np.uint8),
                    "state": rng.randn(2).astype(np.float32),
                    "task": "benchmark",
                }
            )
        dataset.save_episode()
    dataset.finalize()
    return dataset


def _measure_drift(dataset: LeRobotDataset, fps: int, frames_per_episode: list[int]) -> dict:
    episodes = load_episodes(dataset.root)
    num_episodes = len(frames_per_episode)

    chunk_idx = episodes[0][f"videos/{VIDEO_KEY}/chunk_index"]
    file_idx = episodes[0][f"videos/{VIDEO_KEY}/file_index"]
    video_path = dataset.root / dataset.meta.video_path.format(
        video_key=VIDEO_KEY, chunk_index=chunk_idx, file_index=file_idx
    )

    actual_pts = _decode_all_frame_pts(video_path)
    actual_duration = get_video_duration_in_s(video_path)

    boundary_drifts_s: list[float] = []
    cumulative = 0
    single_file = True
    for ep_idx in range(num_episodes):
        # If episodes spilled into multiple files, boundary indexing no longer holds.
        if (
            episodes[ep_idx][f"videos/{VIDEO_KEY}/chunk_index"] != chunk_idx
            or episodes[ep_idx][f"videos/{VIDEO_KEY}/file_index"] != file_idx
        ):
            single_file = False
            break

        if cumulative > 0:
            from_ts = float(episodes[ep_idx][f"videos/{VIDEO_KEY}/from_timestamp"])
            boundary_drifts_s.append(abs(from_ts - actual_pts[cumulative]))
        cumulative += frames_per_episode[ep_idx]

    last_to_ts = float(episodes[num_episodes - 1][f"videos/{VIDEO_KEY}/to_timestamp"])
    duration_drift_s = abs(last_to_ts - actual_duration)

    drifts = np.array(boundary_drifts_s) if boundary_drifts_s else np.array([0.0])
    half_frame_s = 0.5 / fps
    return {
        "num_episodes": num_episodes,
        "num_boundaries": len(boundary_drifts_s),
        "single_file": single_file,
        "total_frames": cumulative,
        "max_drift_s": float(drifts.max()),
        "mean_drift_s": float(drifts.mean()),
        "p99_drift_s": float(np.percentile(drifts, 99)),
        "max_drift_frames": float(drifts.max() * fps),
        "duration_drift_s": duration_drift_s,
        "half_frame_s": half_frame_s,
        "exceeds_half_frame": bool(drifts.max() >= half_frame_s),
    }


def run_config(fps: int, num_episodes: int, min_frames: int, max_frames: int, seed: int, streaming: bool):
    rng = np.random.RandomState(seed)
    frames_per_episode = rng.randint(min_frames, max_frames + 1, size=num_episodes).tolist()
    tmp = Path(tempfile.mkdtemp(prefix="lerobot_drift_bench_"))
    try:
        dataset = _record_dataset(tmp / "dataset", fps, frames_per_episode, streaming)
        return _measure_drift(dataset, fps, frames_per_episode)
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def _print_report(label: str, r: dict) -> None:
    note = "" if r["single_file"] else "  (truncated: episodes spilled to multiple files)"
    print(f"\n=== {label}{note} ===")
    print(f"  episodes / boundaries : {r['num_episodes']} / {r['num_boundaries']}")
    print(f"  total frames          : {r['total_frames']}")
    print(f"  max boundary drift    : {r['max_drift_s']:.3e} s  ({r['max_drift_frames']:.4f} frames)")
    print(f"  mean boundary drift   : {r['mean_drift_s']:.3e} s")
    print(f"  p99 boundary drift    : {r['p99_drift_s']:.3e} s")
    print(f"  total-duration drift  : {r['duration_drift_s']:.3e} s")
    print(f"  half-frame threshold  : {r['half_frame_s']:.3e} s")
    print(f"  exceeds half-frame    : {'YES <-- FAIL' if r['exceeds_half_frame'] else 'no'}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--fps", type=int, default=None, help="Override fps (default: sweep presets).")
    parser.add_argument("--num-episodes", type=int, default=None, help="Override episode count.")
    parser.add_argument("--min-frames", type=int, default=7)
    parser.add_argument("--max-frames", type=int, default=18)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--streaming", action="store_true", help="Use the streaming encoder path.")
    args = parser.parse_args()

    if args.fps is not None or args.num_episodes is not None:
        fps = args.fps or 30
        num_episodes = args.num_episodes or 50
        configs = [(fps, num_episodes)]
    else:
        configs = [(30, 50), (30, 200), (60, 200), (50, 200)]

    for fps, num_episodes in configs:
        r = run_config(fps, num_episodes, args.min_frames, args.max_frames, args.seed, args.streaming)
        label = f"fps={fps}, episodes={num_episodes}, streaming={args.streaming}"
        _print_report(label, r)


if __name__ == "__main__":
    main()