dhb-xr 0.2.1__py3-none-any.whl

examples/integration/test_libero_encoding.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+"""Test full DHB pipeline with real LIBERO-Spatial dataset.
+
+Usage:
+    pixi run python examples/integration/test_libero_encoding.py
+"""
+import os
+import sys
+from pathlib import Path
+
+# Add src to path when running from project root
+sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
+
+import numpy as np
+from dhb_xr.integration.vla.libero import LiberoAdapter
+from dhb_xr.integration.vla.pipeline import DHBVLAPipeline, DHBVLAPipelineConfig
+from dhb_xr.core.types import EncodingMethod, DHBMethod
+
+DATA_DIR = Path(os.environ.get("LIBERO_DATA_DIR", "/home/andypark/Projects/data/libero/libero_spatial"))
+FILE_PATH = DATA_DIR / "pick_up_the_black_bowl_between_the_plate_and_the_ramekin_and_place_it_on_the_plate_demo.hdf5"
+
+print("=== Testing Full DHB Pipeline with LIBERO data ===\n")
+
+# Load episodes using adapter
+adapter = LiberoAdapter()
+episodes = list(adapter.load_dataset(str(FILE_PATH)))[:5]  # First 5 episodes
+print(f"Loaded {len(episodes)} episodes\n")
+
+# Create pipeline
+config = DHBVLAPipelineConfig(
+    dhb_method=DHBMethod.DOUBLE_REFLECTION,
+    method=EncodingMethod.POSITION,
+)
+pipeline = DHBVLAPipeline(config=config)
+print(f"Pipeline config:")
+print(f"  DHB method: {config.dhb_method}")
+print(f"  Encoding method: {config.method}")
+print()
+
+# Process each episode
+for i, episode in enumerate(episodes):
+    pos = episode["positions"]
+    quat = episode["quaternions"]
+    demo_id = episode["metadata"]["demo_id"]
+
+    print(f"--- Episode {i+1}: {demo_id} ---")
+    print(f"  Input: {pos.shape[0]} frames, pos={pos.shape}, quat={quat.shape}")
+
+    try:
+        result = pipeline.encode_trajectory(pos, quat)
+
+        # DHB returns separate linear and angular invariants
+        linear_inv = result["linear_motion_invariants"]
+        angular_inv = result["angular_motion_invariants"]
+        invariants = np.concatenate([linear_inv, angular_inv], axis=1)
+
+        print(f"  Linear invariants: shape={linear_inv.shape}")
+        print(f"  Angular invariants: shape={angular_inv.shape}")
+        print(f"  Combined: shape={invariants.shape}")
+        print(f"  Invariant stats: min={invariants.min():.3f}, max={invariants.max():.3f}, mean={invariants.mean():.3f}")
+
+        # Check for NaN or Inf
+        if np.any(np.isnan(invariants)):
+            print(f"  WARNING: Contains NaN values!")
+        if np.any(np.isinf(invariants)):
+            print(f"  WARNING: Contains Inf values!")
+
+        print()
+    except Exception as e:
+        import traceback
+        print(f"  ERROR: {e}")
+        traceback.print_exc()
+        print()
+
+print("=== SUCCESS: DHB encoding works with LIBERO data! ===")

examples/integration/test_libero_retrieval.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+"""Test motion retrieval with real LIBERO-Spatial dataset.
+
+This demonstrates:
+1. Loading LIBERO episodes using LiberoAdapter
+2. Encoding trajectories with DHB-DR
+3. Building a motion database for retrieval
+4. Querying similar motions across different tasks
+
+Usage:
+    pixi run python examples/integration/test_libero_retrieval.py
+"""
+import os
+import sys
+from pathlib import Path
+
+# Add src to path when running from project root
+sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
+
+from dhb_xr.integration.vla.libero import LiberoAdapter
+from dhb_xr.database.motion_db import MotionDatabase
+
+LIBERO_DIR = Path(os.environ.get("LIBERO_DATA_DIR", "/home/andypark/Projects/data/libero/libero_spatial"))
+
+print("=" * 60)
+print("LIBERO-Spatial Motion Retrieval Demo")
+print("=" * 60)
+print()
+
+# Load episodes from multiple tasks
+adapter = LiberoAdapter()
+all_episodes = []
+task_names = []
+
+hdf5_files = sorted(LIBERO_DIR.glob("*.hdf5"))[:5]  # First 5 tasks
+print(f"Loading episodes from {len(hdf5_files)} tasks...\n")
+
+for hdf5_file in hdf5_files:
+    task_name = hdf5_file.stem.replace("_demo", "")
+    episodes = list(adapter.load_dataset(str(hdf5_file)))[:10]  # 10 episodes per task
+
+    for ep in episodes:
+        ep["metadata"]["task_name"] = task_name
+        ep["metadata"]["file"] = hdf5_file.name
+
+    all_episodes.extend(episodes)
+    task_names.append(task_name)
+    print(f"  {task_name[:60]}...")
+    print(f"    Loaded {len(episodes)} episodes")
+
+print(f"\nTotal: {len(all_episodes)} episodes from {len(task_names)} tasks")
+print()
+
+# Build motion database (encoding happens inside db.add())
+print("Building motion database...")
+db = MotionDatabase(dhb_method="double_reflection")
+
+for i, ep in enumerate(all_episodes):
+    db.add(
+        positions=ep["positions"],
+        quaternions=ep["quaternions"],
+        metadata={**ep["metadata"], "idx": i},
+    )
+
+print(f"  Database contains {len(db.invariants_list)} trajectories")
+print()
+
+# Query similar motions
+print("=" * 60)
+print("Motion Retrieval Results")
+print("=" * 60)
+print()
+
+# Use first episode from each task as query
+for task_idx in range(min(3, len(task_names))):
+    query_idx = task_idx * 10  # First episode of each task
+    query_ep = all_episodes[query_idx]
+    query_task = query_ep["metadata"]["task_name"]
+
+    print(f"Query: {query_task[:50]}...")
+    print(f"  Demo: {query_ep['metadata']['demo_id']}, Frames: {len(query_ep['positions'])}")
+
+    # Find top-6 similar motions (including self)
+    # Use DTW for variable-length trajectories
+    results = db.retrieve(
+        query_positions=query_ep["positions"],
+        query_quaternions=query_ep["quaternions"],
+        k=6,
+        use_dtw=True,
+    )
+
+    print(f"\n  Top-5 Similar Motions:")
+    for j, (invariants, meta, distance) in enumerate(results[1:6]):  # Skip self
+        match_task = meta.get("task_name", "unknown")
+        same_task = "SAME" if match_task == query_task else "DIFF"
+
+        print(f"    {j+1}. [{same_task}] dist={distance:.4f}")
+        print(f"       Task: {match_task[:50]}...")
+        print(f"       Demo: {meta.get('demo_id', 'unknown')}")
+
+    print()
+
+print("=" * 60)
+print("SUCCESS: Motion retrieval working with LIBERO data!")
+print("=" * 60)

examples/motion_database.py

@@ -0,0 +1,88 @@
+"""
+Motion Database Example
+
+This example demonstrates the DHB-XR motion database for trajectory similarity search.
+The database stores trajectories as DHB invariants and supports:
+
+- Efficient similarity search using L2 distance on invariants
+- Optional FAISS for GPU-accelerated nearest neighbor search
+- Dynamic Time Warping (DTW) for variable-length trajectories
+- Metadata storage and retrieval
+
+Applications:
+- Motion planning: find similar trajectories from past experience
+- Imitation learning: retrieve demonstrations by motion similarity
+- Trajectory libraries: organize and search motion collections
+
+The example:
+1. Creates 4 synthetic trajectories with varying randomness
+2. Adds them to the motion database
+3. Queries for the 3 most similar trajectories to a new query
+"""
+
+import numpy as np
+from dhb_xr.database.motion_db import MotionDatabase
+
+
+def run_example():
+    """Run the motion database example."""
+    print("DHB-XR Motion Database Example")
+    print("=" * 35)
+
+    # Initialize database
+    np.random.seed(3)
+    n = 35
+    db = MotionDatabase(dhb_method="double_reflection", use_faiss=False)
+
+    print(f"Creating motion database with {n}-pose trajectories...")
+
+    # Add trajectories to database
+    trajectories_added = []
+    for i in range(4):
+        # Create trajectory with increasing randomness
+        positions = np.cumsum(np.random.randn(n, 3) * (0.01 + i * 0.005), axis=0)
+        quaternions = np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (n, 1))
+
+        metadata = {"id": i, "label": f"traj_{i}", "randomness": 0.01 + i * 0.005}
+        db.add(positions, quaternions, metadata=metadata)
+
+        trajectories_added.append({
+            'start': positions[0],
+            'end': positions[-1],
+            'metadata': metadata
+        })
+
+    print(f"✅ Added {len(trajectories_added)} trajectories to database")
+
+    # Create query trajectory
+    print("\nCreating query trajectory...")
+    q_pos = np.cumsum(np.random.randn(n, 3) * 0.012, axis=0)
+    q_quat = np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (n, 1))
+
+    print(f"Query trajectory: {q_pos[0]} → {q_pos[-1]}")
+
+    # Retrieve similar trajectories
+    print("\nSearching for similar trajectories...")
+    results = db.retrieve(q_pos, q_quat, k=3, use_dtw=False)
+
+    print("Top 3 matches (by L2 invariant distance):")
+    print("-" * 45)
+    for rank, (invariants, metadata, distance) in enumerate(results, 1):
+        traj_info = trajectories_added[metadata['id']]
+        print(f"{rank}. Distance: {distance:.4f}")
+        print(f"   Trajectory: {metadata['label']}")
+        print(f"   Start: {traj_info['start']}")
+        print(f"   End:   {traj_info['end']}")
+        print(f"   Randomness: {metadata['randomness']:.3f}")
+        print()
+
+    return db, results
+
+
+def main():
+    """Main function for command-line execution."""
+    run_example()
+
+
+if __name__ == "__main__":
+    main()

examples/trajectory_adaptation.py

@@ -0,0 +1,85 @@
+"""
+Trajectory Adaptation Example
+
+This example demonstrates trajectory adaptation: retargeting a demonstration
+trajectory to start and end at different poses while preserving the motion shape.
+
+The process:
+1. Create a demo trajectory (40 poses with random walk)
+2. Define new start and goal poses (offset from original)
+3. Use DHB encoding/decoding to adapt the trajectory to the new boundary conditions
+4. Show that the adapted trajectory starts/ends at the desired poses
+
+This is the foundation for robust trajectory execution under pose perturbations.
+"""
+
+import numpy as np
+from dhb_xr.optimization.casadi_solver import generate_trajectory
+from dhb_xr.core.types import DHBMethod
+
+
+def run_example():
+    """Run the trajectory adaptation example."""
+    print("DHB-XR Trajectory Adaptation Example")
+    print("=" * 40)
+
+    # Create demo trajectory
+    np.random.seed(0)
+    n = 40
+    positions = np.cumsum(np.random.randn(n, 3) * 0.02, axis=0)
+    quaternions = np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (n, 1))
+
+    print(f"Demo trajectory: {n} poses")
+    print(f"Original start: {positions[0]}")
+    print(f"Original goal:  {positions[-1]}")
+
+    # Define new boundary conditions
+    pose_init = {"position": positions[0].copy(), "quaternion": quaternions[0].copy()}
+    pose_goal = {
+        "position": positions[-1].copy() + np.array([0.1, 0.1, 0.0]),
+        "quaternion": quaternions[-1].copy()
+    }
+
+    print("
+New goal (offset by [0.1, 0.1, 0]):"
+    print(f"Target goal: {pose_goal['position']}")
+
+    # Adapt trajectory
+    print("\nAdapting trajectory...")
+    result = generate_trajectory(
+        positions, quaternions,
+        pose_init, pose_goal,
+        traj_length=30,
+        smoothing=False,
+        dhb_method=DHBMethod.DOUBLE_REFLECTION,
+    )
+
+    adapted_positions = result["adapted_pos_data"]
+    print(f"Adapted trajectory: {len(adapted_positions)} poses")
+    print(f"Adapted start: {adapted_positions[0]}")
+    print(f"Adapted goal:  {adapted_positions[-1]}")
+
+    # Check boundary condition satisfaction
+    start_error = np.linalg.norm(adapted_positions[0] - pose_init["position"])
+    goal_error = np.linalg.norm(adapted_positions[-1] - pose_goal["position"])
+
+    print("
+Boundary condition errors:")
+    print(".2e")
+    print(".2e")
+
+    if start_error < 1e-10 and goal_error < 1e-10:
+        print("✅ Perfect boundary condition satisfaction")
+    else:
+        print("⚠️  Boundary conditions not perfectly satisfied")
+
+    return result
+
+
+def main():
+    """Main function for command-line execution."""
+    run_example()
+
+
+if __name__ == "__main__":
+    main()

examples/vla_tokenization.py

@@ -0,0 +1,107 @@
+"""
+VLA Tokenization Example (DHB-Token)
+
+This example demonstrates VQ-VAE tokenization of DHB invariants for discrete
+action representations suitable for VLA (Vision-Language-Action) models.
+
+DHB-Token converts continuous DHB invariants into discrete tokens:
+1. Encode trajectory to DHB-DR invariants (8 values per timestep)
+2. Use VQ-VAE to quantize to discrete tokens
+3. Reconstruct invariants from tokens
+
+Benefits for VLA:
+- Discrete action space (like decision transformers)
+- Compression: many timesteps → few tokens
+- Autoregressive generation capabilities
+- Robust to noise and perturbations
+
+Requires: PyTorch (`pip install dhb_xr[gpu]`)
+"""
+
+import numpy as np
+
+
+def run_example():
+    """Run the VLA tokenization example."""
+    print("DHB-XR VLA Tokenization (DHB-Token) Example")
+    print("=" * 45)
+
+    try:
+        import torch
+        from dhb_xr.encoder.dhb_dr import encode_dhb_dr
+        from dhb_xr.core.types import DHBMethod, EncodingMethod
+        from dhb_xr.tokenization.vqvae import DHBTokenizer
+    except ImportError as e:
+        print("❌ PyTorch required for tokenization")
+        print(f"Install with: pip install dhb_xr[gpu]")
+        print(f"Error: {e}")
+        return None
+
+    # Create synthetic trajectory
+    np.random.seed(2)
+    n = 40
+    positions = np.cumsum(np.random.randn(n, 3) * 0.01, axis=0)
+    quaternions = np.tile(np.array([1.0, 0.0, 0.0, 0.0]), (n, 1))
+
+    print(f"Trajectory: {n} poses")
+    print(f"Positions shape: {positions.shape}")
+
+    # Encode to DHB invariants
+    print("\nEncoding to DHB-DR invariants...")
+    out = encode_dhb_dr(
+        positions,
+        quaternions,
+        method=EncodingMethod.POSITION,
+        use_default_initial_frames=True,
+        dhb_method=DHBMethod.DOUBLE_REFLECTION,
+    )
+
+    # Combine linear and angular invariants
+    invariants = np.concatenate([
+        out["linear_motion_invariants"],
+        out["angular_motion_invariants"],
+    ], axis=1).astype(np.float32)
+
+    print(f"Combined invariants shape: {invariants.shape}")
+    print(f"Values per timestep: {invariants.shape[1]} (4 linear + 4 angular)")
+
+    # Convert to PyTorch tensor and add batch dimension
+    inv_batch = torch.from_numpy(invariants).unsqueeze(0)  # (T, 8) -> (1, T, 8)
+    print(f"Batch shape: {inv_batch.shape}")
+
+    # Tokenize with VQ-VAE
+    print("\nTokenizing with VQ-VAE...")
+    model = DHBTokenizer(invariant_dim=8, latent_dim=16, codebook_size=64)
+    indices, reconstructed, z, z_q = model(inv_batch)
+
+    print(f"✅ Tokenization complete")
+    print(f"Token indices shape: {indices.shape}")
+    print(f"Reconstructed shape: {reconstructed.shape}")
+
+    # Calculate reconstruction quality
+    rec_loss = torch.nn.functional.mse_loss(reconstructed, inv_batch)
+    compression_ratio = invariants.size / indices.numel()
+
+    print("
+Results:")
+    print(".4f")
+    print(".1f")
+
+    # Show some token examples
+    print(f"\nSample tokens: {indices[0, :10].tolist()}")
+
+    return {
+        'invariants': invariants,
+        'tokens': indices,
+        'reconstructed': reconstructed,
+        'loss': rec_loss.item()
+    }
+
+
+def main():
+    """Main function for command-line execution."""
+    run_example()
+
+
+if __name__ == "__main__":
+    main()

notebooks/__init__.py

@@ -0,0 +1,24 @@
+"""
+DHB-XR Tutorial Notebooks
+
+This subpackage contains Jupyter notebook tutorials for DHB-XR:
+
+Quick Start Tutorials:
+- tutorial_dhb_basics.ipynb: Basic encoding/decoding concepts
+- tutorial_invariance_demo.ipynb: SE(3) invariance demonstration
+- tutorial_umi_xdof_motion_analysis.ipynb: Real-world UMI dataset analysis
+
+Advanced Topics:
+- demo_dhb_features_and_use_cases.ipynb: Comprehensive feature overview
+- tutorial_vla_integration.ipynb: VLA integration with LIBERO/RoboCASA
+
+Research & Development:
+- manuscript_figures.ipynb: Figures for research manuscripts
+
+Run notebooks with:
+    jupyter notebook notebooks/
+    # or
+    pixi run notebook
+"""
+
+__version__ = "0.2.0"