File size: 2,828 Bytes
adc02fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | #!/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())
|