scripts/generate_rlbench_lattice.py

#!/usr/bin/env python
"""
Phase B: Generate RLBench CIL dataset (second benchmark)
Adapts ManiSkill lattice generation to RLBench

This demonstrates method generality beyond ManiSkill - critical for A* paper.
"""
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))

# TODO: Implement RLBench backend or use alternative second benchmark
# For now, this is a placeholder showing the structure

def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Generate RLBench CIL dataset (Phase B: Second Benchmark)"
    )
    parser.add_argument(
        "--tasks",
        type=str,
        default="reach_target,take_lid_off_saucepan,open_drawer",
        help="Comma-separated RLBench task names"
    )
    parser.add_argument("--num-groups", type=int, default=1000)
    parser.add_argument("--k", type=int, default=16)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=0)

    args = parser.parse_args(argv)

    print("=" * 60)
    print("Phase B: RLBench CIL Generation")
    print("=" * 60)
    print(f"Tasks: {args.tasks}")
    print(f"Groups per task: {args.num_groups}")
    print(f"K (interventions): {args.k}")
    print(f"Output: {args.out}")
    print()

    # Check if RLBench is available
    try:
        import rlbench  # noqa: F401
        print("✅ RLBench package found")
    except ImportError:
        print("❌ RLBench not installed")
        print()
        print("ALTERNATIVE: Use Meta-World instead (faster to integrate)")
        print()
        print("To install Meta-World:")
        print("  pip install metaworld")
        print()
        print("Meta-World has 50 manipulation tasks, easier API than RLBench")
        print("Good second benchmark choice for A* paper")
        return 1

    # TODO: Implement RLBench CIL generation
    # Structure should mirror generate_maniskill_lattice.py:
    # 1. Load RLBench demonstrations
    # 2. For each demo, sample N states
    # 3. For each state, serialize and generate K intervention candidates
    # 4. Restore state and execute each candidate
    # 5. Record outcomes as CIL records
    # 6. Save to sharded JSONL format

    print("⚠️  RLBench integration not yet implemented")
    print()
    print("RECOMMENDED PATH:")
    print("1. Use Meta-World instead (pip install metaworld)")
    print("2. Adapt this script for Meta-World API")
    print("3. Meta-World is faster to integrate and widely recognized")
    print()
    print("See: scripts/generate_metaworld_lattice.py (to be created)")

    return 0

if __name__ == "__main__":
    sys.exit(main())