scripts/visualize_verb_object_samples.py

#!/usr/bin/env python3
"""
Sample a small number of VerbObjectColor-v1 trajectories (motion planning) and save short videos.

  python scripts/visualize_verb_object_samples.py --verb push --color green --shape cup -n 2
  python scripts/visualize_verb_object_samples.py --verb rotate --shape cube --vis

From repo root; use CPU backends if you hit Vulkan/mplib segfaults.
"""
from __future__ import annotations

import argparse
import os
import sys

script_dir = os.path.dirname(os.path.abspath(__file__))
root = os.path.dirname(script_dir)
if root not in sys.path:
    sys.path.insert(0, root)

import mani_skill.envs  # noqa: F401

from scripts.run_verb_color_shape_motion_planning import run_verb_color_shape_motion_planning


def main():
    p = argparse.ArgumentParser(description="Few-sample MP demos for VerbObjectColor-v1")
    p.add_argument("--verb", type=str, default="push")
    p.add_argument("--color", type=str, default="red")
    p.add_argument("--shape", type=str, default="cube")
    p.add_argument("-n", type=int, default=2, help="Number of trajectories")
    p.add_argument("--obj-xy", type=float, nargs=2, default=None)
    p.add_argument("--record-dir", type=str, default="demos/visual_samples")
    p.add_argument("--traj-name", type=str, default="viz_sample")
    p.add_argument("--distractor-max", type=int, default=3)
    p.add_argument("-b", "--sim-backend", type=str, default="cpu", choices=["cpu", "gpu"])
    p.add_argument("--render-backend", type=str, default="cpu", choices=["cpu", "gpu"])
    p.add_argument("--vis", action="store_true", help="Open viewer (slow)")
    p.add_argument("--no-video", action="store_true")
    args = p.parse_args()

    run_verb_color_shape_motion_planning(
        verb=args.verb,
        color=args.color,
        shape=args.shape,
        obj_xy=list(args.obj_xy) if args.obj_xy is not None else None,
        num_traj=args.n,
        only_count_success=False,
        sim_backend=args.sim_backend,
        render_backend=args.render_backend,
        save_video=not args.no_video,
        vis=args.vis,
        record_dir=args.record_dir,
        traj_name=args.traj_name,
        distractor_max=args.distractor_max,
    )


if __name__ == "__main__":
    main()