scripts/render_maniskill_observations.py

#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from dovla_cil.generation.maniskill_render import (  # noqa: E402
    ManiSkillRenderConfig,
    render_maniskill_observations,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Render RGB observations from persisted ManiSkill CIL states."
    )
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--render-backend", default="cpu")
    parser.add_argument("--image-quality", type=int, default=90)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--overwrite", action="store_true")
    args = parser.parse_args(argv)
    result = render_maniskill_observations(
        ManiSkillRenderConfig(
            dataset_dir=args.dataset,
            render_backend=args.render_backend,
            image_quality=args.image_quality,
            seed=args.seed,
            overwrite=args.overwrite,
        )
    )
    print(json.dumps(result, indent=2))
    return 0


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