dhb-xr 0.2.1__py3-none-any.whl

examples/integration/libero_simulation_demo.py

@@ -0,0 +1,286 @@
+#!/usr/bin/env python
+"""LIBERO Simulation Demo with DHB-XR.
+
+This script demonstrates:
+1. Setting up a LIBERO environment
+2. Loading a demo trajectory from HDF5
+3. Encoding/decoding with DHB invariants
+4. Replaying the trajectory in simulation
+
+Requirements:
+    conda create -n libero python=3.10
+    conda activate libero
+    pip install robosuite mujoco h5py numpy
+    git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git
+    cd LIBERO && pip install -e .
+    pip install dhb_xr
+
+Usage:
+    python examples/integration/libero_simulation_demo.py --task_id 0 --render
+
+Author: Andy Park
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+# Add src to path for development
+sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
+
+import numpy as np
+
+# Check for LIBERO installation
+try:
+    from libero.libero import benchmark, get_libero_path
+    from libero.libero.envs import OffScreenRenderEnv
+    HAS_LIBERO = True
+except ImportError:
+    HAS_LIBERO = False
+    print("LIBERO not installed. Install with:")
+    print("  git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git")
+    print("  cd LIBERO && pip install -e .")
+
+# Check for robosuite
+try:
+    import robosuite
+    HAS_ROBOSUITE = True
+except ImportError:
+    HAS_ROBOSUITE = False
+    print("robosuite not installed. Install with: pip install robosuite mujoco")
+
+
+def load_demo_from_hdf5(dataset_path: str, demo_id: int = 0):
+    """Load a demo trajectory from LIBERO HDF5 dataset."""
+    import h5py
+
+    with h5py.File(dataset_path, "r") as f:
+        demo_key = f"demo_{demo_id}"
+        if demo_key not in f["data"]:
+            raise ValueError(f"Demo {demo_id} not found in {dataset_path}")
+
+        demo = f["data"][demo_key]
+
+        # Extract actions
+        actions = np.array(demo["actions"])
+
+        # Extract end-effector pose from robot_states
+        robot_states = np.array(demo["robot_states"])
+        ee_pos = robot_states[:, 2:5]  # columns 2-4
+        ee_quat_wxyz = robot_states[:, 5:9]  # columns 5-8 (w, x, y, z)
+
+        # Convert to (x, y, z, w) format
+        ee_quat = ee_quat_wxyz[:, [1, 2, 3, 0]]
+
+        return {
+            "actions": actions,
+            "positions": ee_pos,
+            "quaternions": ee_quat,
+            "num_frames": len(actions),
+        }
+
+
+def encode_demo_with_dhb(positions: np.ndarray, quaternions: np.ndarray):
+    """Encode trajectory with DHB-DR."""
+    from dhb_xr.encoder.dhb_dr import encode_dhb_dr
+    from dhb_xr.core.types import EncodingMethod, DHBMethod
+
+    result = encode_dhb_dr(
+        positions, quaternions,
+        method=EncodingMethod.POSITION,
+        dhb_method=DHBMethod.DOUBLE_REFLECTION,
+    )
+
+    return {
+        "linear_invariants": result["linear_motion_invariants"],
+        "angular_invariants": result["angular_motion_invariants"],
+        "initial_pose": result["initial_pose"],
+    }
+
+
+def decode_demo_with_dhb(
+    linear_inv: np.ndarray,
+    angular_inv: np.ndarray,
+    initial_pose: dict,
+):
+    """Decode DHB invariants back to trajectory."""
+    from dhb_xr.decoder.dhb_dr import decode_dhb_dr
+    from dhb_xr.core.types import EncodingMethod, DHBMethod
+
+    result = decode_dhb_dr(
+        linear_inv, angular_inv,
+        initial_pose,
+        method=EncodingMethod.POSITION,
+        dhb_method=DHBMethod.DOUBLE_REFLECTION,
+        drop_padded=True,
+    )
+
+    return {
+        "positions": result["positions"],
+        "quaternions": result["quaternions"],
+    }
+
+
+def run_libero_demo(task_id: int = 0, demo_id: int = 0, render: bool = False):
+    """Run a LIBERO task with a demo trajectory."""
+
+    if not HAS_LIBERO or not HAS_ROBOSUITE:
+        print("\nCannot run simulation without LIBERO and robosuite.")
+        print("Running DHB encoding/decoding demo only...\n")
+
+        # Use our downloaded dataset for DHB demo
+        dataset_dir = Path(os.environ.get(
+            "LIBERO_DATA_DIR",
+            "/home/andypark/Projects/data/libero/libero_spatial"
+        ))
+
+        if not dataset_dir.exists():
+            print(f"Dataset not found at {dataset_dir}")
+            print("Download with:")
+            print("  wget -O libero_spatial.zip 'https://utexas.box.com/shared/static/04k94hyizn4huhbv5sz4ev9p2h1p6s7f.zip'")
+            return
+
+        # Find first HDF5 file
+        hdf5_files = list(dataset_dir.glob("*.hdf5"))
+        if not hdf5_files:
+            print(f"No HDF5 files found in {dataset_dir}")
+            return
+
+        dataset_path = str(hdf5_files[0])
+        print(f"Loading demo from: {dataset_path}")
+
+        demo = load_demo_from_hdf5(dataset_path, demo_id)
+        print(f"Loaded demo with {demo['num_frames']} frames")
+        print(f"  Positions shape: {demo['positions'].shape}")
+        print(f"  Quaternions shape: {demo['quaternions'].shape}")
+
+        # Encode with DHB
+        print("\nEncoding with DHB-DR...")
+        encoded = encode_demo_with_dhb(demo["positions"], demo["quaternions"])
+        print(f"  Linear invariants: {encoded['linear_invariants'].shape}")
+        print(f"  Angular invariants: {encoded['angular_invariants'].shape}")
+
+        # Decode back
+        print("\nDecoding back to trajectory...")
+        decoded = decode_demo_with_dhb(
+            encoded["linear_invariants"],
+            encoded["angular_invariants"],
+            encoded["initial_pose"],
+        )
+        print(f"  Decoded positions: {decoded['positions'].shape}")
+        print(f"  Decoded quaternions: {decoded['quaternions'].shape}")
+
+        # Compute reconstruction error
+        pos_error = np.linalg.norm(demo["positions"] - decoded["positions"][:len(demo["positions"])], axis=1)
+        print(f"\nReconstruction error:")
+        print(f"  Position: mean={pos_error.mean():.6f}m, max={pos_error.max():.6f}m")
+
+        print("\n=== DHB encoding/decoding successful! ===")
+        print("\nTo run full simulation, install LIBERO:")
+        print("  conda create -n libero python=3.10")
+        print("  conda activate libero")
+        print("  pip install robosuite mujoco")
+        print("  git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git")
+        print("  cd LIBERO && pip install -e .")
+        return
+
+    # ========== Full LIBERO simulation ==========
+    print(f"Setting up LIBERO task {task_id}...")
+
+    # Get benchmark and task
+    benchmark_dict = benchmark.get_benchmark_dict()
+    task_suite = benchmark_dict["libero_spatial"]()
+
+    task = task_suite.get_task(task_id)
+    print(f"Task: {task.name}")
+    print(f"Description: {task.language}")
+
+    # Get BDDL file for environment setup
+    task_bddl_file = os.path.join(
+        get_libero_path("bddl_files"),
+        task.problem_folder,
+        task.bddl_file
+    )
+
+    # Create environment
+    env_args = {
+        "bddl_file_name": task_bddl_file,
+        "camera_heights": 128,
+        "camera_widths": 128,
+    }
+
+    if render:
+        env_args["has_renderer"] = True
+        env_args["has_offscreen_renderer"] = False
+
+    env = OffScreenRenderEnv(**env_args)
+    env.seed(0)
+
+    # Get initial states for benchmarking
+    init_states = task_suite.get_task_init_states(task_id)
+    init_state_id = 0
+
+    # Reset and set initial state
+    env.reset()
+    env.set_init_state(init_states[init_state_id])
+
+    # Load demo actions from dataset
+    dataset_path = os.path.join(
+        get_libero_path("datasets"),
+        "libero_spatial",
+        f"{task.name}_demo.hdf5"
+    )
+
+    print(f"Loading demo from: {dataset_path}")
+    demo = load_demo_from_hdf5(dataset_path, demo_id)
+
+    # Encode with DHB
+    print("Encoding demo with DHB-DR...")
+    encoded = encode_demo_with_dhb(demo["positions"], demo["quaternions"])
+
+    print(f"Running {demo['num_frames']} steps in simulation...")
+
+    # Execute demo actions
+    for step in range(demo["num_frames"]):
+        action = demo["actions"][step]
+        obs, reward, done, info = env.step(action)
+
+        if render:
+            env.render()
+
+        if done:
+            print(f"Task completed at step {step}!")
+            break
+
+    # Check final reward
+    if reward > 0:
+        print("SUCCESS: Task completed!")
+    else:
+        print("Task not completed within demo length.")
+
+    env.close()
+    print("\nSimulation finished.")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="LIBERO + DHB-XR Demo")
+    parser.add_argument("--task_id", type=int, default=0, help="LIBERO task ID")
+    parser.add_argument("--demo_id", type=int, default=0, help="Demo ID within task")
+    parser.add_argument("--render", action="store_true", help="Render simulation")
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("LIBERO Simulation Demo with DHB-XR")
+    print("=" * 60)
+    print()
+
+    run_libero_demo(
+        task_id=args.task_id,
+        demo_id=args.demo_id,
+        render=args.render,
+    )
+
+
+if __name__ == "__main__":
+    main()