Refactored code to allow for multiple build environments

DEPENDENCY_CONFLICT.md

@@ -0,0 +1,87 @@
+# Dependency Conflict: OpenPI vs OpenVLA
+
+## Problem
+
+OpenPI and OpenVLA cannot coexist in the same Python environment due to incompatible CUDA/PyTorch requirements:
+
+### OpenVLA Requirements
+- `torch==2.2.0` (CUDA 12.1)
+- `nvidia-cudnn-cu12==8.9.2.26`
+- `transformers==4.40.1`
+- `draccus==0.8.0`
+
+### OpenPI Requirements  
+- `torch>=2.7.0` (CUDA 12.8)
+- `nvidia-cudnn-cu12>=9.1.1`
+- `transformers==4.48.1`
+- `draccus>=0.10.0`
+
+### The Core Issue
+When both are installed together:
+1. OpenVLA downgrades torch from 2.9.1 → 2.2.0
+2. OpenVLA downgrades cuDNN from 9.10 → 8.9
+3. OpenPI's JAX components are compiled against cuDNN 9.1+ but runtime loads cuDNN 8.9/9.0
+4. Result: **"Loaded runtime CuDNN library: 9.0.0 but source was compiled with: 9.1.1"**
+
+This causes OpenPI initialization to fail with `ImportError: initialization failed`.
+
+## Solution Options
+
+### Option 1: Separate Containers (Recommended for Production)
+Run OpenPI and OpenVLA in separate Docker containers or Hugging Face Spaces:
+- **OpenPI Space**: Uses current setup with torch>=2.7.0
+- **OpenVLA Space**: Separate space with torch==2.2.0
+
+### Option 2: Conditional Installation
+Install only one model at a time based on environment variable:
+```bash
+if [ "$MODEL" = "openvla" ]; then
+    pip install git+https://github.com/openvla/openvla.git
+else
+    pip install git+https://github.com/tan7271/OpenPiRoboEval.git
+fi
+```
+
+### Option 3: Virtual Environments (Local Development)
+Use separate conda/venv environments:
+```bash
+# OpenPI environment
+conda create -n openpi python=3.11
+conda activate openpi
+# ... install OpenPI deps
+
+# OpenVLA environment  
+conda create -n openvla python=3.11
+conda activate openvla
+# ... install OpenVLA deps
+```
+
+## Current Implementation
+
+This repository currently supports **OpenPI only**. OpenVLA has been removed from:
+- `requirements.txt` (removed peft, timm, accelerate that were added for OpenVLA)
+- `setup.sh` (removed OpenVLA git installation)
+- `app.py` (removed OpenVLA from MODEL_REGISTRY)
+
+## Re-enabling OpenVLA
+
+To create a separate OpenVLA-only deployment:
+
+1. Revert to commit before OpenPI was added
+2. Use these versions in `requirements.txt`:
+```
+torch==2.2.0
+torchvision==0.17.0
+torchaudio==2.2.0
+transformers==4.40.1
+```
+3. Add back OpenVLA installation to `setup.sh`
+4. Keep only OpenVLA in `MODEL_REGISTRY`
+
+## References
+
+- OpenPI repo: https://github.com/tan7271/OpenPiRoboEval
+- OpenVLA repo: https://github.com/openvla/openvla
+- Related issue: CUDA library version mismatches causing initialization failures
+
+

README.md

@@ -1,5 +1,5 @@
 ---
-title: Pi0 Inference on RoboEval Tasks
+title: Robot Policy Inference on RoboEval Tasks
 emoji: 🤖
 colorFrom: blue
 colorTo: purple
@@ -12,17 +12,33 @@ license: mit
 python_version: "3.11"
 ---
 
-# Pi0 Inference on RoboEval Tasks
+# Robot Policy Inference on RoboEval Tasks
 
-A Hugging Face Space for running Pi0 bimanual manipulation policy inference on various robot tasks from the RoboEval benchmark.
+A Hugging Face Space for running robot manipulation policy inference on various tasks from the RoboEval benchmark. Supports **OpenPI** (Pi0 bimanual policy) and **OpenVLA** (vision-language-action) backends.
 
 ## 🚀 Features
 
+- **Multiple Model Backends**: Switch between OpenPI and OpenVLA using environment variables
 - **Interactive Gradio Interface**: Easy-to-use web interface for running inference
 - **Multiple Tasks**: Support for 20+ bimanual manipulation tasks
 - **Real-time Video Output**: View robot execution videos immediately after inference
 - **Customizable Parameters**: Adjust max steps, FPS, and task instructions
 - **GPU Acceleration**: Runs on T4 GPU for fast inference
+- **Dynamic Model Detection**: Interface adapts based on which backend is installed
+
+## 🔀 Switching Between OpenPI and OpenVLA
+
+This Space supports both OpenPI and OpenVLA, but they **cannot run simultaneously** due to dependency conflicts. Choose your backend using the `MODEL_BACKEND` environment variable:
+
+### Quick Setup
+
+1. Go to **Settings → Variables** in your Space
+2. Add environment variable:
+   - **Name**: `MODEL_BACKEND`
+   - **Value**: `openpi` (default) or `openvla`
+3. Save and rebuild
+
+See [SWITCHING_MODELS.md](./SWITCHING_MODELS.md) for detailed instructions and technical explanation.
 
 ## 📋 Available Tasks
 

SWITCHING_MODELS.md

@@ -0,0 +1,124 @@
+# Switching Between OpenPI and OpenVLA
+
+This Space supports both OpenPI and OpenVLA backends, but they **cannot run simultaneously** due to dependency conflicts. You can switch between them using an environment variable.
+
+## How to Switch Models
+
+### Option 1: Using Hugging Face Space Settings (Recommended)
+
+1. Go to your Space settings: **Settings → Variables**
+2. Add a new **Environment Variable**:
+   - **Name**: `MODEL_BACKEND`
+   - **Value**: `openpi` or `openvla`
+3. Click **Save**
+4. The Space will rebuild with the selected backend
+
+### Option 2: Local Development
+
+Set the environment variable before running:
+
+```bash
+# For OpenPI (default)
+export MODEL_BACKEND=openpi
+bash setup.sh
+python app.py
+
+# For OpenVLA
+export MODEL_BACKEND=openvla
+bash setup.sh
+python app.py
+```
+
+## What Happens During Build
+
+The `setup.sh` script checks the `MODEL_BACKEND` variable:
+
+- **`MODEL_BACKEND=openpi`** (default):
+  - Installs PyTorch 2.9+ with cuDNN 9.1+
+  - Installs lerobot and OpenPI from git
+  - OpenPI appears in the model dropdown
+
+- **`MODEL_BACKEND=openvla`**:
+  - Installs PyTorch 2.2.0 with cuDNN 8.9
+  - Installs OpenVLA from git
+  - OpenVLA appears in the model dropdown
+
+## Dynamic Model Detection
+
+The app automatically detects which backend is installed:
+
+```python
+# In app.py
+def _populate_model_registry():
+    try:
+        import openpi
+        # Register OpenPI
+    except ImportError:
+        pass
+    
+    try:
+        import openvla
+        # Register OpenVLA  
+    except ImportError:
+        pass
+```
+
+The Gradio interface will only show the models that are actually available.
+
+## Why Can't Both Run Together?
+
+**Dependency Conflict Summary:**
+
+| Package | OpenPI Requirement | OpenVLA Requirement | Conflict |
+|---------|-------------------|---------------------|----------|
+| torch | >=2.7.0 | ==2.2.0 | ❌ Incompatible |
+| nvidia-cudnn-cu12 | >=9.1.1 | ==8.9.2.26 | ❌ Incompatible |
+| transformers | ==4.48.1 | ==4.40.1 | ❌ Incompatible |
+| draccus | ==0.10.0 | ==0.8.0 | ❌ Incompatible |
+
+When both are installed together, OpenVLA downgrades PyTorch and cuDNN, causing OpenPI's JAX components to fail with:
+```
+Loaded runtime CuDNN library: 9.0.0 but source was compiled with: 9.1.1
+```
+
+See [`DEPENDENCY_CONFLICT.md`](./DEPENDENCY_CONFLICT.md) for detailed technical explanation.
+
+## Testing Locally
+
+```bash
+# Test OpenPI
+export MODEL_BACKEND=openpi
+bash setup.sh
+python app.py
+
+# Clean environment
+pip uninstall -y openpi lerobot torch torchvision torchaudio
+
+# Test OpenVLA
+export MODEL_BACKEND=openvla  
+bash setup.sh
+python app.py
+```
+
+## Default Behavior
+
+If `MODEL_BACKEND` is not set, the Space defaults to **OpenPI**.
+
+## Troubleshooting
+
+### Space shows "No model backends available"
+- Check build logs for installation errors
+- Verify `MODEL_BACKEND` is set correctly
+- Ensure `GH_TOKEN` secret is configured for private repos
+
+### Wrong model appears after switching
+- Space caching may cause old builds to persist
+- Try: **Settings → Factory Reboot**
+- Or: Change any other setting to force a full rebuild
+
+### Want to use both models?
+Create two separate Spaces:
+- `your-space-openpi` with `MODEL_BACKEND=openpi`
+- `your-space-openvla` with `MODEL_BACKEND=openvla`
+
+

app.py

@@ -1,444 +1,148 @@
 """
-Hugging Face Space for Pi0 Inference on RoboEval Tasks
+Hugging Face Space for Robot Policy Inference on RoboEval Tasks
 
-This Gradio app allows users to run Pi0 model inference on bimanual robot tasks
-and view the resulting execution videos.
+This Gradio app allows users to run model inference (OpenPI, OpenVLA) on bimanual robot tasks
+and view the resulting execution videos. Models run in isolated conda environments.
 """
 
 import os
-import tempfile
-import copy
-import numpy as np
+import json
+import atexit
 import dataclasses
-from pathlib import Path
-from typing import Callable, Dict, List, Optional, Tuple
+from dataclasses import asdict
+from typing import Callable, Dict, Optional, Tuple
 import gradio as gr
 import subprocess
 import sys
+import datetime
 
-# --- Headless defaults (set BEFORE mujoco/roboeval imports) ---
+# --- Headless defaults ---
 os.environ.setdefault("MUJOCO_GL", "egl")
 os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
 os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp")
 
-# Note: Dependencies are installed via setup.sh before the app starts
-# This keeps the app code clean and separates installation logic
+# Note: Model dependencies are installed in separate conda environments via setup.sh
+# This app runs in the base environment and dispatches to subprocess workers
+
+print(f"===== Application Startup at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} =====\n")
 
-# Run setup if dependencies aren't installed
-def check_and_install_dependencies():
-    """Check if dependencies are installed, run setup if not."""
-    dependencies_ok = True
+# Verify environments exist on startup
+def verify_environments():
+    """Check that conda environments exist"""
+    result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True)
     
-    try:
-        import roboeval
-        print("✓ roboeval imported")
-    except ImportError as e:
-        print(f"✗ roboeval import failed: {e}")
-        dependencies_ok = False
+    has_openpi = "openpi_env" in result.stdout
+    has_openvla = "openvla_env" in result.stdout
     
-    try:
-        import lerobot
-        print("✓ lerobot imported")
-    except ImportError as e:
-        print(f"✗ lerobot import failed: {e}")
-        dependencies_ok = False
+    print("Environment check:")
+    print(f"  {'✓' if has_openpi else '✗'} openpi_env")
+    print(f"  {'✓' if has_openvla else '✗'} openvla_env (optional)")
     
-    try:
-        import openpi
-        print("✓ openpi imported")
-    except ImportError as e:
-        print(f"✗ openpi import failed: {e}")
-        dependencies_ok = False
+    if not has_openpi:
+        raise RuntimeError("openpi_env not found. Check setup.sh logs.")
     
-    # If core dependencies are missing, run setup
-    if not dependencies_ok:
-        print("\n" + "="*60)
-        print("INSTALLING MISSING DEPENDENCIES")
-        print("="*60)
-        print("Running setup.sh to install roboeval, lerobot, and openpi...")
-        
-        import subprocess
-        import os
-        
-        setup_path = os.path.join(os.path.dirname(__file__), "setup.sh")
-        result = subprocess.run(
-            ["bash", setup_path],
-            cwd=os.path.dirname(__file__),
-            capture_output=False,
-            text=True
-        )
-        
-        if result.returncode != 0:
-            print(f"Setup script failed with return code {result.returncode}")
-            raise RuntimeError("Setup script failed to install dependencies")
-        
-        print("\n" + "="*60)
-        print("SETUP COMPLETE - Verifying installations...")
-        print("="*60)
-        
-        # Verify installations
-        try:
-            import roboeval
-            print("✓ roboeval installed successfully")
-        except ImportError as e:
-            print(f"✗ roboeval still not available: {e}")
-            
-        try:
-            import lerobot
-            print("✓ lerobot installed successfully")
-        except ImportError as e:
-            print(f"✗ lerobot still not available: {e}")
-            
-        try:
-            import openpi
-            print("✓ openpi installed successfully")
-        except ImportError as e:
-            print(f"✗ openpi still not available: {e}")
-    
-    return True
+    return has_openpi, has_openvla
 
-import datetime
-print(f"===== Application Startup at {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} =====\n")
-check_and_install_dependencies()
-
-# --- OpenPI (local inference) ---
-try:
-    from openpi.training import config as _config
-    from openpi.policies import policy_config as _policy_config
-    OPENPI_AVAILABLE = True
-    print("OpenPI imported successfully")
-except ImportError as e:
-    print(f"Error: OpenPI import failed after installation: {e}")
-    
-    # All dependencies should be in requirements.txt now
-    print(f"OpenPI import failed. Check that all dependencies are properly installed.")
-    OPENPI_AVAILABLE = False
-
-# --- RoboEval imports ---
-from roboeval.action_modes import JointPositionActionMode
-from roboeval.utils.observation_config import ObservationConfig, CameraConfig
-from roboeval.robots.configs.panda import BimanualPanda
-from roboeval.roboeval_env import CONTROL_FREQUENCY_MAX
-
-# Import all environment classes
-from roboeval.envs.manipulation import (
-    CubeHandover, CubeHandoverOrientation, CubeHandoverPosition,
-    CubeHandoverPositionAndOrientation, VerticalCubeHandover,
-    StackTwoBlocks, StackTwoBlocksOrientation, StackTwoBlocksPosition,
-    StackTwoBlocksPositionAndOrientation
-)
-from roboeval.envs.lift_pot import (
-    LiftPot, LiftPotOrientation, LiftPotPosition, LiftPotPositionAndOrientation,
-)
-from roboeval.envs.lift_tray import (
-    LiftTray, DragOverAndLiftTray, LiftTrayOrientation, LiftTrayPosition, LiftTrayPositionAndOrientation,
-)
-from roboeval.envs.pack_objects import (
-    PackBox, PackBoxOrientation, PackBoxPosition, PackBoxPositionAndOrientation,
-)
-from roboeval.envs.stack_books import (
-    PickSingleBookFromTable, PickSingleBookFromTableOrientation,
-    PickSingleBookFromTablePosition, PickSingleBookFromTablePositionAndOrientation,
-    StackSingleBookShelf, StackSingleBookShelfPosition, StackSingleBookShelfPositionAndOrientation,
-)
-from roboeval.envs.rotate_utility_objects import (
-    RotateValve, RotateValveObstacle, RotateValvePosition, RotateValvePositionAndOrientation,
-)
-
-# --- Video ---
-from moviepy.editor import VideoClip
+HAS_OPENPI, HAS_OPENVLA = verify_environments()
 
 # ---------------------- Environment Registry ----------------------
+# Task names for UI dropdown (workers have their own registry)
 _ENV_CLASSES = {
-    "CubeHandover": (CubeHandover, "handover the rod from one hand to the other hand"),
-    "CubeHandoverOrientation": (CubeHandoverOrientation, "handover the rod from one hand to the other hand"),
-    "CubeHandoverPosition": (CubeHandoverPosition, "handover the rod from one hand to the other hand"),
-    "CubeHandoverPositionOrientation": (CubeHandoverPositionAndOrientation, "handover the rod from one hand to the other hand"),
-    "CubeHandoverVertical": (VerticalCubeHandover, "handover the rod from one hand to the other hand"),
-
-    "LiftPot": (LiftPot, "lift the pot by the handles"),
-    "LiftPotOrientation": (LiftPotOrientation, "lift the pot by the handles"),
-    "LiftPotPosition": (LiftPotPosition, "lift the pot by the handles"),
-    "LiftPotPositionOrientation": (LiftPotPositionAndOrientation, "lift the pot by the handles"),
-
-    "LiftTray": (LiftTray, "lift the tray"),
-    "LiftTrayDrag": (DragOverAndLiftTray, "lift the tray"),
-    "LiftTrayOrientation": (LiftTrayOrientation, "lift the tray"),
-    "LiftTrayPosition": (LiftTrayPosition, "lift the tray"),
-    "LiftTrayPositionOrientation": (LiftTrayPositionAndOrientation, "lift the tray"),
-
-    "PackBox": (PackBox, "close the box"),
-    "PackBoxOrientation": (PackBoxOrientation, "close the box"),
-    "PackBoxPosition": (PackBoxPosition, "close the box"),
-    "PackBoxPositionOrientation": (PackBoxPositionAndOrientation, "close the box"),
-
-    "PickSingleBookFromTable": (PickSingleBookFromTable, "pick up the book from the table"),
-    "PickSingleBookFromTableOrientation": (PickSingleBookFromTableOrientation, "pick up the book from the table"),
-    "PickSingleBookFromTablePosition": (PickSingleBookFromTablePosition, "pick up the book from the table"),
-    "PickSingleBookFromTablePositionOrientation": (PickSingleBookFromTablePositionAndOrientation, "pick up the book from the table"),
-
-    "RotateValve": (RotateValve, "rotate the valve counter clockwise"),
-    "RotateValveObstacle": (RotateValveObstacle, "rotate the valve counter clockwise"),
-    "RotateValvePosition": (RotateValvePosition, "rotate the valve counter clockwise"),
-    "RotateValvePositionOrientation": (RotateValvePositionAndOrientation, "rotate the valve counter clockwise"),
-
-    "StackSingleBookShelf": (StackSingleBookShelf, "put the book on the table onto the shelf"),
-    "StackSingleBookShelfPosition": (StackSingleBookShelfPosition, "put the book on the table onto the shelf"),
-    "StackSingleBookShelfPositionOrientation": (StackSingleBookShelfPositionAndOrientation, "put the book on the table onto the shelf"),
-
-    "StackTwoBlocks": (StackTwoBlocks, "stack the two cubes"),
-    "StackTwoBlocksOrientation": (StackTwoBlocksOrientation, "stack the two cubes"),
-    "StackTwoBlocksPosition": (StackTwoBlocksPosition, "stack the two cubes"),
-    "StackTwoBlocksPositionOrientation": (StackTwoBlocksPositionAndOrientation, "stack the two cubes")
+    "CubeHandover": "handover the rod from one hand to the other hand",
+    "CubeHandoverOrientation": "handover the rod from one hand to the other hand",
+    "CubeHandoverPosition": "handover the rod from one hand to the other hand",
+    "CubeHandoverPositionOrientation": "handover the rod from one hand to the other hand",
+    "CubeHandoverVertical": "handover the rod from one hand to the other hand",
+    "LiftPot": "lift the pot by the handles",
+    "LiftPotOrientation": "lift the pot by the handles",
+    "LiftPotPosition": "lift the pot by the handles",
+    "LiftPotPositionOrientation": "lift the pot by the handles",
+    "LiftTray": "lift the tray",
+    "LiftTrayDrag": "lift the tray",
+    "LiftTrayOrientation": "lift the tray",
+    "LiftTrayPosition": "lift the tray",
+    "LiftTrayPositionOrientation": "lift the tray",
+    "PackBox": "close the box",
+    "PackBoxOrientation": "close the box",
+    "PackBoxPosition": "close the box",
+    "PackBoxPositionOrientation": "close the box",
+    "PickSingleBookFromTable": "pick up the book from the table",
+    "PickSingleBookFromTableOrientation": "pick up the book from the table",
+    "PickSingleBookFromTablePosition": "pick up the book from the table",
+    "PickSingleBookFromTablePositionOrientation": "pick up the book from the table",
+    "RotateValve": "rotate the valve counter clockwise",
+    "RotateValveObstacle": "rotate the valve counter clockwise",
+    "RotateValvePosition": "rotate the valve counter clockwise",
+    "RotateValvePositionOrientation": "rotate the valve counter clockwise",
+    "StackSingleBookShelf": "put the book on the table onto the shelf",
+    "StackSingleBookShelfPosition": "put the book on the table onto the shelf",
+    "StackSingleBookShelfPositionOrientation": "put the book on the table onto the shelf",
+    "StackTwoBlocks": "stack the two cubes",
+    "StackTwoBlocksOrientation": "stack the two cubes",
+    "StackTwoBlocksPosition": "stack the two cubes",
+    "StackTwoBlocksPositionOrientation": "stack the two cubes"
 }
 
 # ---------------------- Configuration ----------------------
-DEFAULT_DEVICE = "cuda:0" if os.path.exists("/dev/nvidia0") else "cpu"
-DEFAULT_DOWNSAMPLE_RATE = 25
 DEFAULT_MAX_STEPS = 200
 DEFAULT_FPS = 25
 
-# Check GPU availability and print diagnostics
-def check_gpu_status():
-    """Check and print GPU availability."""
-    import jax
-    print("\n" + "="*60)
-    print("GPU DIAGNOSTICS")
-    print("="*60)
-    
-    # Check JAX devices
-    devices = jax.devices()
-    print(f"JAX devices: {devices}")
-    print(f"JAX default backend: {jax.default_backend()}")
-    
-    # Check if GPU is available
-    gpu_devices = [d for d in devices if d.platform == 'gpu']
-    if gpu_devices:
-        print(f"✅ GPU available! Found {len(gpu_devices)} GPU device(s)")
-        for i, device in enumerate(gpu_devices):
-            print(f"   GPU {i}: {device}")
-    else:
-        print(f"❌ No GPU found. Running on: {jax.default_backend()}")
-    
-    # Check CUDA
-    try:
-        import torch
-        if torch.cuda.is_available():
-            print(f"✅ PyTorch CUDA available: {torch.cuda.get_device_name(0)}")
-            print(f"   CUDA version: {torch.version.cuda}")
-            print(f"   GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
-        else:
-            print("❌ PyTorch CUDA not available")
-    except Exception as e:
-        print(f"⚠️  Could not check PyTorch CUDA: {e}")
-    
-    print("="*60 + "\n")
-    return len(gpu_devices) > 0
-
-# Run GPU check at startup
-GPU_AVAILABLE = check_gpu_status()
-
-# Global policy cache to avoid reloading
-_POLICY_CACHE = {}
-
-def clear_gpu_memory():
-    """Clear GPU memory and policy cache."""
-    global _POLICY_CACHE
-    
-    # Clear the policy cache
-    _POLICY_CACHE.clear()
-    
-    # Force JAX to clear GPU memory
-    try:
-        import jax
-        import gc
-        
-        # Clear JAX compilation caches (for JAX >= 0.4.36)
-        try:
-            jax.clear_caches()
-        except AttributeError:
-            # Fallback for older JAX versions
-            try:
-                jax.clear_backends()
-            except AttributeError:
-                pass  # Neither method available, rely on gc
-        
-        # Force Python garbage collection
-        gc.collect()
-        
-        print("GPU memory cleared successfully")
-    except Exception as e:
-        print(f"Warning: Could not fully clear GPU memory: {e}")
+# ---------------------- Subprocess Worker Management ----------------------
+# Global: persistent subprocess pool
+_INFERENCE_WORKERS: Dict[str, subprocess.Popen] = {
+    "openpi": None,
+    "openvla": None
+}
 
-# ---------------------- OpenPI Helpers ----------------------
-def get_checkpoint_path(task_name: str, ckpt_path: Optional[str] = None) -> str:
-    """
-    Return a local path to the checkpoint for the given task. If `ckpt_path` is provided,
-    it is returned verbatim. Otherwise, download only the files under:
-        repo: tan7271/pi0_base_checkpoints (model repo)
-        subdir: {task_name}_testing/{step}/
-    We prefer step=2999 if present, else the numerically largest available step.
 
-    This version avoids snapshot_download entirely to dodge the "0 files" issue.
+def get_inference_worker(model_key: str) -> subprocess.Popen:
     """
-    if ckpt_path:
-        return ckpt_path
-
-    from huggingface_hub import HfApi, hf_hub_download
-
-    repo_id = "tan7271/pi0_base_checkpoints"
-    revision = "main"
-    base_dir = f"{task_name}_testing"
-    cache_dir = os.path.expanduser("~/.cache/roboeval/pi0_checkpoints")
-
-    api = HfApi()
-    try:
-        all_files: List[str] = api.list_repo_files(
-            repo_id=repo_id, revision=revision, repo_type="model"
-        )
-    except Exception as e:
-        raise RuntimeError(f"Could not list files for {repo_id}@{revision}: {e}")
-
-    # Find available numeric steps under {task}_testing/
-    steps = sorted({
-        int(p.split("/")[1])
-        for p in all_files
-        if p.startswith(base_dir + "/") and len(p.split("/")) >= 3 and p.split("/")[1].isdigit()
-    })
-    if not steps:
-        nearby = [p for p in all_files if base_dir in p][:10]
-        raise FileNotFoundError(
-            f"No files found under '{base_dir}/' in {repo_id}@{revision}. "
-            f"Example paths I do see: {nearby}"
-        )
-
-    chosen_step = 2999 if 2999 in steps else steps[-1]
-    subdir = f"{base_dir}/{chosen_step}"
-
-    print(
-        f"Downloading checkpoint for {task_name} directly via hf_hub_download "
-        f"(repo={repo_id}, subdir={subdir})..."
-    )
-
-    # We only need these parts; if you want rollouts, drop the filter below.
-    needed_roots = (
-        f"{subdir}/_CHECKPOINT_METADATA",
-        f"{subdir}/assets/",
-        f"{subdir}/params/",
-        f"{subdir}/train_state/",
-    )
-    wanted = [
-        p for p in all_files
-        if p == f"{subdir}/_CHECKPOINT_METADATA"
-        or any(p.startswith(root) for root in needed_roots[1:])
-    ]
-
-    # If the filtered list is empty (unexpected), grab the entire subdir.
-    if not wanted:
-        wanted = [p for p in all_files if p.startswith(subdir + "/")]
-        if not wanted:
-            raise FileNotFoundError(
-                f"Repo listing shows no files under '{subdir}/'. "
-                f"Steps seen: {steps}"
-            )
-
-    manual_root = os.path.join(cache_dir, "manual")
-    os.makedirs(manual_root, exist_ok=True)
-
-    # Download every file we want into a local mirror of the repo layout.
-    for relpath in wanted:
-        hf_hub_download(
-            repo_id=repo_id,
-            filename=relpath,
-            revision=revision,
-            repo_type="model",
-            local_dir=manual_root,
-            local_dir_use_symlinks=True,  # saves space on shared filesystems
-        )
-
-    manual_ckpt_dir = os.path.join(manual_root, subdir)
-
-    # Basic sanity: ensure the directory exists and isn't empty
-    def _nonempty_dir(path: str) -> bool:
-        return os.path.isdir(path) and any(True for _ in os.scandir(path))
-
-    if not _nonempty_dir(manual_ckpt_dir):
-        try:
-            siblings = [e.name for e in os.scandir(os.path.dirname(manual_ckpt_dir))]
-        except Exception:
-            siblings = []
-        raise FileNotFoundError(
-            f"Downloaded files, but '{manual_ckpt_dir}' is missing/empty.\n"
-            f"Siblings present: {siblings}\n"
-            f"(repo_id={repo_id}, subdir={subdir})"
-        )
-
-    return manual_ckpt_dir
-
-
-def load_pi0_base_bimanual_droid(task_name: str, ckpt_path: str):
-    """Load Pi0 policy model for the given task."""
-    if not OPENPI_AVAILABLE:
-        raise RuntimeError("OpenPI is not available. Cannot load Pi0 model.")
-    
-    # Get checkpoint path (download from HF if needed)
-    checkpoint_path = get_checkpoint_path(task_name, ckpt_path)
+    Get or create persistent inference worker subprocess.
     
-    cache_key = f"{task_name}:{checkpoint_path}"
-    if cache_key in _POLICY_CACHE:
-        return _POLICY_CACHE[cache_key]
+    Workers stay alive to keep models loaded in memory (fast subsequent calls).
+    """
+    global _INFERENCE_WORKERS
     
-    # Clear old policies from cache to free GPU memory for new task
-    if len(_POLICY_CACHE) > 0:
-        print(f"Clearing {len(_POLICY_CACHE)} cached model(s) to free GPU memory...")
-        clear_gpu_memory()
+    # Check if environment exists
+    env_name = f"{model_key}_env"
+    result = subprocess.run(["conda", "env", "list"], capture_output=True, text=True)
+    if env_name not in result.stdout:
+        raise RuntimeError(f"Environment {env_name} not found. Check setup.sh logs.")
     
-    cfg = _config.get_config("pi0_base_bimanual_droid_finetune")
-    bimanual_assets = _config.AssetsConfig(
-        assets_dir=f"{checkpoint_path}/assets/",
-        asset_id=f"tan7271/{task_name}",
-    )
-    cfg = dataclasses.replace(cfg, data=dataclasses.replace(cfg.data, assets=bimanual_assets))
-    policy = _policy_config.create_trained_policy(cfg, checkpoint_path)
+    if _INFERENCE_WORKERS[model_key] is None or _INFERENCE_WORKERS[model_key].poll() is not None:
+        # Start new worker
+        script_name = f"inference_{model_key}.py"
+        
+        print(f"Starting {model_key} worker in {env_name}...")
+        
+        proc = subprocess.Popen(
+            ["conda", "run", "-n", env_name, "python", script_name],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,  # Line buffered
+        )
+        
+        _INFERENCE_WORKERS[model_key] = proc
+        print(f"✓ {model_key} worker started (PID: {proc.pid})")
     
-    _POLICY_CACHE[cache_key] = policy
-    return policy
-
-
-def make_openpi_example_from_roboeval(obs_dict: dict, prompt: str) -> dict:
-    """Convert RoboEval observation to OpenPI format."""
-    obs = obs_dict[0] if isinstance(obs_dict, (tuple, list)) else obs_dict
-    example = {"prompt": prompt}
-
-    # Cameras (CHW→HWC)
-    exterior_chw = obs["rgb_head"]
-    left_wrist_chw = obs["rgb_left_wrist"]
-    right_wrist_chw = obs["rgb_right_wrist"]
-    example["observation/exterior_image_1_left"] = np.moveaxis(exterior_chw, 0, -1)
-    example["observation/wrist_image_left"] = np.moveaxis(left_wrist_chw, 0, -1)
-    example["observation/wrist_image_right"] = np.moveaxis(right_wrist_chw, 0, -1)
-
-    # Joints and grippers
-    prop = np.asarray(obs["proprioception"], dtype=np.float32).reshape(-1)
-    example["observation/joint_position"] = prop
+    return _INFERENCE_WORKERS[model_key]
 
-    grip = np.asarray(obs["proprioception_grippers"], dtype=np.float32).reshape(-1)[:2]
-    example["observation/gripper_position"] = grip
-
-    return example
-
-
-def map_policy_action_to_env_abs(action_vec: np.ndarray, env) -> np.ndarray:
-    """Map policy action to environment action."""
-    a = np.asarray(action_vec, dtype=np.float32).reshape(-1)
-    if a.shape[0] != 16:
-        raise ValueError(f"Expected (16,), got {a.shape}.")
-    return a
 
+def cleanup_workers():
+    """Terminate worker subprocesses on shutdown"""
+    for model_key, proc in _INFERENCE_WORKERS.items():
+        if proc and proc.poll() is None:
+            print(f"Terminating {model_key} worker...")
+            proc.terminate()
+            try:
+                proc.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                proc.kill()
+                proc.wait()
 
-def _clip_to_space(env, action: np.ndarray) -> np.ndarray:
-    """Safety: clip to env action space."""
-    return np.clip(action, env.action_space.low, env.action_space.high)
+atexit.register(cleanup_workers)
 
 
 @dataclasses.dataclass
@@ -450,6 +154,16 @@ class InferenceRequest:
     max_steps: int
     fps: int
     progress: gr.Progress
+    
+    def to_dict(self) -> Dict:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "task_name": self.task_name,
+            "checkpoint_path": self.checkpoint_path or "",
+            "custom_instruction": self.custom_instruction,
+            "max_steps": self.max_steps,
+            "fps": self.fps,
+        }
 
 
 @dataclasses.dataclass
@@ -460,224 +174,78 @@ class ModelDefinition:
     run_inference: Callable[[InferenceRequest], Tuple[Optional[str], str]]
 
 
-# ---------------------- Video Helpers ----------------------
-def save_frames_to_video(frames, output_path: str, fps: int = 25) -> str:
-    """Save rollout frames to a video file."""
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    
-    # Convert frames to proper format
-    frames = np.array(frames)
-    if frames.ndim == 4 and frames.shape[1] == 3:  # (T, C, H, W)
-        frames = np.moveaxis(frames, 1, -1)  # → (T, H, W, C)
-    
-    duration = len(frames) / fps
-    clip = VideoClip(make_frame=lambda t: frames[min(int(t * fps), len(frames) - 1)], duration=duration)
-    clip.write_videofile(output_path, fps=fps, codec="libx264", logger=None)
-    
-    return output_path
-
-
-# ---------------------- RoboEval Environment Setup ----------------------
-def setup_env(task_name: str, downsample_rate: int = 25):
-    """Setup RoboEval environment for the given task."""
-    cameras = [
-        CameraConfig(name="head", rgb=True, depth=False, resolution=(256, 256)),
-        CameraConfig(name="left_wrist", rgb=True, depth=False, resolution=(256, 256)),
-        CameraConfig(name="right_wrist", rgb=True, depth=False, resolution=(256, 256)),
-    ]
-    
-    Env, prompt = _ENV_CLASSES[task_name]
-    env = Env(
-        action_mode=JointPositionActionMode(
-            floating_base=True,
-            absolute=False,
-            block_until_reached=False,
-            ee=False,
-            floating_dofs=[],
-        ),
-        observation_config=ObservationConfig(cameras=cameras, proprioception=True),
-        render_mode="rgb_array",
-        robot_cls=BimanualPanda,
-        control_frequency=CONTROL_FREQUENCY_MAX // downsample_rate,
-    )
-    
-    return env, prompt
-
-
-def _unpack_obs(obs_or_tuple):
-    """Unpack observation if it's a tuple."""
-    return obs_or_tuple[0] if isinstance(obs_or_tuple, (tuple, list)) else obs_or_tuple
-
-
-# ---------------------- Inference Loop ----------------------
-def run_inference_loop(
-    env,
-    policy,
-    instruction: str,
-    max_steps: int = 200,
-    open_loop_horizon: int = 10,
-):
-    """Run inference loop for one episode."""
-    obs = env.reset()
-    successes = 0
-    images_env = []
-    chunk = None
-    i_in_chunk = 0
-
-    for step_idx in range(max_steps):
-        cur_obs = _unpack_obs(obs)
-        
-        # Collect environment render
-        images_env.append(copy.deepcopy(env.render()))
-        
-        # Request new action chunk when needed
-        if chunk is None or i_in_chunk >= open_loop_horizon:
-            example = make_openpi_example_from_roboeval(cur_obs, instruction)
-            out = policy.infer(example)
-            chunk = out["actions"]
-            i_in_chunk = 0
-
-        # Take next action from cached chunk
-        a_vec = chunk[i_in_chunk]
-        i_in_chunk += 1
-
-        env_action = map_policy_action_to_env_abs(a_vec, env)
-        env_action = _clip_to_space(env, env_action)
-
-        obs, reward, terminated, truncated, info = env.step(env_action)
-        successes += int(reward > 0)
-
-        if terminated or truncated:
-            break
-
-    stats = {"steps": step_idx + 1, "success_signal": successes}
-    return stats, images_env
-
-
-# ---------------------- Main Inference Function ----------------------
+# ---------------------- Main Inference Functions ----------------------
 def run_pi0_inference(request: InferenceRequest) -> Tuple[Optional[str], str]:
-    """
-    Main function to run Pi0 inference.
-    
-    Returns:
-        Tuple of (video_path, status_message)
-    """
+    """Dispatch OpenPI inference to subprocess"""
     try:
-        task_name = request.task_name
-        checkpoint_path = request.checkpoint_path
-        custom_instruction = request.custom_instruction
-        max_steps = int(request.max_steps)
-        fps = int(request.fps)
-        progress = request.progress
-
-        progress(0, desc="Loading model and environment...")
-        
-        # Check GPU status
-        import jax
-        gpu_info = ""
-        devices = jax.devices()
-        gpu_devices = [d for d in devices if d.platform == 'gpu']
-        if gpu_devices:
-            gpu_info = f"🎮 **GPU**: {len(gpu_devices)} GPU(s) detected - {gpu_devices[0]}\n"
-        else:
-            gpu_info = f"⚠️ **GPU**: Not detected! Running on {jax.default_backend()}\n"
+        request.progress(0, desc="Starting OpenPI worker...")
+        worker = get_inference_worker("openpi")
         
-        # Check if OpenPI is available
-        if not OPENPI_AVAILABLE:
-            return None, gpu_info + f"❌ **OpenPI not available**\n\nOpenPI is required for Pi0 model inference but is not installed. Please check the build logs for installation errors."
+        request.progress(0.1, desc="Sending inference request...")
+        # Send request
+        request_dict = request.to_dict()
+        request_json = json.dumps(request_dict)
+        worker.stdin.write(request_json + "\n")
+        worker.stdin.flush()
         
-        # Validate task
-        if task_name not in _ENV_CLASSES:
-            return None, f"❌ Unknown task: {task_name}"
+        request.progress(0.2, desc="Waiting for inference result...")
+        # Read result
+        result_line = worker.stdout.readline()
+        if not result_line:
+            return None, "❌ Worker process ended unexpectedly"
         
-        # Load policy
-        progress(0.2, desc="Loading Pi0 policy...")
-        policy = load_pi0_base_bimanual_droid(task_name, checkpoint_path)
+        result = json.loads(result_line)
         
-        # Setup environment
-        progress(0.4, desc="Setting up environment...")
-        env, default_prompt = setup_env(task_name, downsample_rate=DEFAULT_DOWNSAMPLE_RATE)
-        instruction = custom_instruction if custom_instruction else default_prompt
+        request.progress(1.0, desc="Complete!")
         
-        # Run inference
-        progress(0.5, desc="Running inference...")
-        stats, images_env = run_inference_loop(
-            env, policy, instruction, max_steps=max_steps
-        )
-        
-        # Save video
-        progress(0.8, desc="Saving video...")
-        video_path = os.path.join(tempfile.gettempdir(), f"pi0_rollout_{task_name}.mp4")
-        save_frames_to_video(images_env, video_path, fps=fps)
-        
-        # Cleanup
-        env.close()
+        if result["success"]:
+            return result["video_path"], result["status_message"]
+        else:
+            error_msg = f"❌ OpenPI Error: {result.get('error', 'Unknown error')}\n\n{result.get('status_message', '')}"
+            return None, error_msg
+            
+    except Exception as e:
+        import traceback
+        return None, f"❌ Worker communication error: {str(e)}\n\n{traceback.format_exc()}"
+
+
+def run_openvla_inference(request: InferenceRequest) -> Tuple[Optional[str], str]:
+    """Dispatch OpenVLA inference to subprocess"""
+    try:
+        request.progress(0, desc="Starting OpenVLA worker...")
+        worker = get_inference_worker("openvla")
         
-        # Clear GPU memory after inference to prevent OOM on next run
-        import gc
-        gc.collect()
+        request.progress(0.1, desc="Sending inference request...")
+        # Send request
+        request_dict = request.to_dict()
+        request_json = json.dumps(request_dict)
+        worker.stdin.write(request_json + "\n")
+        worker.stdin.flush()
         
-        progress(1.0, desc="Complete!")
+        request.progress(0.2, desc="Waiting for inference result...")
+        # Read result
+        result_line = worker.stdout.readline()
+        if not result_line:
+            return None, "❌ Worker process ended unexpectedly"
         
-        status = gpu_info + f"✅ **Inference Complete!**\n\n"
-        status += f"- **Task**: {task_name}\n"
-        status += f"- **Steps**: {stats['steps']}\n"
-        status += f"- **Success Signal**: {stats['success_signal']}\n"
-        status += f"- **Instruction**: {instruction}\n"
+        result = json.loads(result_line)
         
-        return video_path, status
+        request.progress(1.0, desc="Complete!")
         
+        if result["success"]:
+            return result["video_path"], result["status_message"]
+        else:
+            error_msg = f"❌ OpenVLA Error: {result.get('error', 'Unknown error')}\n\n{result.get('status_message', '')}"
+            return None, error_msg
+            
     except Exception as e:
         import traceback
-        
-        # Check if it's an out of memory error
-        if "Out of memory" in str(e) or "RESOURCE_EXHAUSTED" in str(e):
-            error_msg = f"""❌ **Out of Memory Error**
-
-The model is too large for the current hardware configuration.
-
-**Pi0 Model Requirements:**
-- Minimum: 8 GB GPU memory
-- Recommended: 16+ GB GPU memory
-
-**Solutions:**
-1. **Upgrade this Space to use a GPU** (Settings → Hardware → T4 small or better)
-2. Use a smaller/quantized checkpoint
-3. Contact the Space owner to enable GPU hardware
-
-**Technical Details:**
-```
-{str(e)}
-```
-"""
-        else:
-            error_msg = f"❌ **Error during inference:**\n\n```\n{str(e)}\n\n{traceback.format_exc()}\n```"
-        
-        return None, error_msg
-
-
-def run_openvla_inference(request: InferenceRequest) -> Tuple[Optional[str], str]:
-    """
-    Placeholder for OpenVLA backend integration.
-    
-    Currently returns a descriptive message until the OpenVLA runtime is wired up.
-    """
-    status = (
-        "⚠️ **OpenVLA integration is not yet available in this Space.**\n\n"
-        "The frontend is model-aware, so you can wire in the backend by implementing "
-        "`run_openvla_inference` to load checkpoints and execute rollouts.\n\n"
-        "Requested configuration:\n"
-        f"- Task: {request.task_name}\n"
-        f"- Checkpoint: {request.checkpoint_path or 'auto'}\n"
-        f"- Steps: {request.max_steps}\n"
-        f"- FPS: {request.fps}\n"
-    )
-    return None, status
+        return None, f"❌ Worker communication error: {str(e)}\n\n{traceback.format_exc()}"
 
 
-# Registry of supported models (UI order follows this definition)
+# Registry of supported models (populated dynamically based on available environments)
 MODEL_REGISTRY: Dict[str, ModelDefinition] = {
-    "pi0_openpi": ModelDefinition(
+    "openpi": ModelDefinition(
         label="Pi0 Base (OpenPI)",
         description=(
             "Runs the Pi0 bimanual policy using the OpenPI runtime. "
@@ -686,16 +254,21 @@ MODEL_REGISTRY: Dict[str, ModelDefinition] = {
         ),
         run_inference=run_pi0_inference,
     ),
-    "openvla": ModelDefinition(
+}
+
+# Add OpenVLA only if environment exists
+if HAS_OPENVLA:
+    MODEL_REGISTRY["openvla"] = ModelDefinition(
         label="OpenVLA",
         description=(
             "Runs the OpenVLA (Open Vision-Language-Action) policy. "
             "OpenVLA is a vision-language-action model for robot manipulation tasks. "
-            "Provide a checkpoint path or leave empty to use default OpenVLA checkpoints."
+            "**Checkpoint path is required** - provide a path to an OpenVLA checkpoint directory."
         ),
         run_inference=run_openvla_inference,
-    ),
-}
+    )
+else:
+    print("ℹ OpenVLA environment not found - OpenVLA model will not be available")
 
 
 def _format_model_info(model_key: str) -> str:
@@ -744,12 +317,16 @@ def create_gradio_interface():
         gr.Markdown("""
         # 🤖 Robot Policy Inference on RoboEval Tasks
         
-        Choose a supported model backend (starting with Pi0 via OpenPI) and run it on RoboEval tasks to watch the generated execution video.
+        Choose a supported model backend and run it on RoboEval tasks to watch the generated execution video.
+        
+        **Architecture**: Models run in isolated conda environments to avoid dependency conflicts. The first inference with each model may take 30-60 seconds to load the model, but subsequent inferences are fast.
         
         ⚠️ **Hardware Requirements:** This Space requires a GPU with at least 8GB memory.
         If you see "Out of Memory" errors, upgrade the Space hardware in Settings → Hardware → T4 small.
         
-        **Note**: Leave the checkpoint path empty to use the model's default retrieval logic (Pi0 fetches from `tan7271/pi0_base_checkpoints`), or provide a custom local path.
+        **Checkpoint Paths**:
+        - **OpenPI**: Leave empty to auto-download from `tan7271/pi0_base_checkpoints`, or provide a custom path
+        - **OpenVLA**: **Required** - provide a path to an OpenVLA checkpoint directory
         """)
         
         with gr.Row():
@@ -846,6 +423,10 @@ def create_gradio_interface():
         - **Stack Books**: Manipulate books on tables and shelves
         - And many more variations with position/orientation constraints
         
+        ### Model Switching
+        
+        You can switch between OpenPI and OpenVLA models instantly. Each model runs in its own isolated environment with optimized dependencies. The first inference with each model loads it into memory (30-60s), but subsequent inferences are fast.
+        
         ### GPU Usage
         
         This Space uses a T4 GPU (~$0.60/hour). It auto-sleeps after 10 minutes of inactivity to minimize costs.

eval_openVLA.py

@@ -0,0 +1,515 @@
+"""
+OpenVLA Evaluation Script for Bimanual Robot Tasks
+
+This script evaluates OpenVLA models on bimanual robot manipulation tasks,
+with support for both demonstration replay and model inference modes.
+
+Usage Examples:
+    # Model inference mode:
+    python 4_eval_openvla.py --ckpt_path /path/to/model/checkpoint
+
+    # Demonstration replay mode:
+    python 4_eval_openvla.py --ckpt_path /path/to/model/checkpoint \
+                             --use_demos --dataset_path /path/to/demo/dataset
+
+    # Custom configuration:
+    python 4_eval_openvla.py --ckpt_path /path/to/model/checkpoint \
+                             --instruction "pick up the red object" \
+                             --num_episodes 10 --max_steps 300 \
+                             --output_dir /path/to/output/videos
+
+Required Arguments:
+    --ckpt_path: Path to the OpenVLA model checkpoint directory
+
+Optional Arguments:
+    --dataset_path: Path to demonstration dataset (required if --use_demos is set)
+    --use_demos: Use demonstration replay instead of model inference
+    --instruction: Task instruction for the robot
+    --device: Device for model inference (default: cuda:0)
+    --downsample_rate: Control frequency downsampling factor (default: 25)
+    --max_steps: Maximum steps per episode (default: 200)
+    --num_episodes: Number of episodes to run (default: 5)
+    --fps: FPS for output videos (default: 5)
+    --output_dir: Output directory for videos (default: checkpoint directory)
+"""
+
+import argparse
+import copy
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import (
+    AutoConfig,
+    AutoImageProcessor,
+    AutoModelForVision2Seq,
+    AutoProcessor,
+)
+
+if not os.environ.get("DISPLAY"):
+    os.environ.setdefault("MUJOCO_GL", "egl")
+    os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
+    os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp")
+
+from prismatic.extern.hf.configuration_prismatic import OpenVLAConfig
+from prismatic.extern.hf.modeling_prismatic import OpenVLAForActionPrediction
+from prismatic.extern.hf.processing_prismatic import PrismaticImageProcessor, PrismaticProcessor
+
+from roboeval.action_modes import JointPositionActionMode
+from roboeval.demonstrations.demo_converter import DemoConverter
+from roboeval.demonstrations.demo_store import DemoStore
+from roboeval.demonstrations.utils import Metadata
+from roboeval.envs.lift_pot import LiftPot
+from roboeval.roboeval_env import CONTROL_FREQUENCY_MAX
+from roboeval.robots.configs.panda import BimanualPanda
+from roboeval.utils.observation_config import CameraConfig, ObservationConfig
+
+try:
+    from moviepy.editor import VideoClip
+except ImportError:
+    raise ImportError("Install moviepy for video preview: pip install moviepy pygame")
+
+
+# Configuration constants
+DEFAULT_DEVICE = "cuda:0"
+DEFAULT_DOWNSAMPLE_RATE = 25
+DEFAULT_MAX_STEPS = 200
+DEFAULT_NUM_EPISODES = 5
+DEFAULT_FPS = 25
+CAMERA_RESOLUTION = (256, 256)
+MIN_REWARD_THRESHOLD = 0.25
+
+
+@dataclass
+class EvaluationConfig:
+    """Configuration for OpenVLA evaluation."""
+    ckpt_path: str = None
+    dataset_path: str = None
+    use_demos: bool = False
+    instruction: str = "reach the red sphere with the left hand and the green sphere with the right hand"
+    device: str = DEFAULT_DEVICE
+    downsample_rate: int = DEFAULT_DOWNSAMPLE_RATE
+    max_steps: int = DEFAULT_MAX_STEPS
+    num_episodes: int = DEFAULT_NUM_EPISODES
+    fps: int = DEFAULT_FPS
+    output_dir: str = None  # Optional output directory for videos
+
+    @classmethod
+    def from_args(cls, args: argparse.Namespace) -> 'EvaluationConfig':
+        """Create config from command line arguments."""
+        return cls(
+            ckpt_path=args.ckpt_path,
+            dataset_path=args.dataset_path,
+            use_demos=args.use_demos,
+            instruction=args.instruction,
+            device=args.device,
+            downsample_rate=args.downsample_rate,
+            max_steps=args.max_steps,
+            num_episodes=args.num_episodes,
+            fps=args.fps,
+            output_dir=args.output_dir
+        )
+
+    def validate(self) -> None:
+        """Validate configuration parameters."""
+        if self.ckpt_path is None:
+            raise ValueError("ckpt_path must be specified")
+        if self.use_demos and self.dataset_path is None:
+            raise ValueError("dataset_path must be specified when use_demos=True")
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Evaluate OpenVLA models on bimanual robot tasks",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    
+    # Required arguments
+    parser.add_argument(
+        "--ckpt_path", 
+        type=str, 
+        required=True,
+        help="Path to the OpenVLA model checkpoint directory"
+    )
+    
+    # Optional arguments
+    parser.add_argument(
+        "--dataset_path", 
+        type=str,
+        help="Path to the demonstration dataset (required if --use_demos is set)"
+    )
+    
+    parser.add_argument(
+        "--use_demos", 
+        action="store_true",
+        help="Use demonstration replay instead of model inference"
+    )
+    
+    parser.add_argument(
+        "--instruction", 
+        type=str,
+        default="reach the red sphere with the left hand and the green sphere with the right hand",
+        help="Task instruction for the robot"
+    )
+    
+    parser.add_argument(
+        "--device", 
+        type=str,
+        default=DEFAULT_DEVICE,
+        help="Device for model inference"
+    )
+    
+    parser.add_argument(
+        "--downsample_rate", 
+        type=int,
+        default=DEFAULT_DOWNSAMPLE_RATE,
+        help="Control frequency downsampling factor"
+    )
+    
+    parser.add_argument(
+        "--max_steps", 
+        type=int,
+        default=DEFAULT_MAX_STEPS,
+        help="Maximum number of steps per episode"
+    )
+    
+    parser.add_argument(
+        "--num_episodes", 
+        type=int,
+        default=DEFAULT_NUM_EPISODES,
+        help="Number of episodes to run"
+    )
+    
+    parser.add_argument(
+        "--fps", 
+        type=int,
+        default=DEFAULT_FPS,
+        help="FPS for output videos"
+    )
+    
+    parser.add_argument(
+        "--output_dir", 
+        type=str,
+        help="Output directory for videos (defaults to checkpoint directory)"
+    )
+    
+    return parser.parse_args()
+
+
+def register_openvla() -> None:
+    """Register OpenVLA components with the Transformers library."""
+    AutoConfig.register("openvla", OpenVLAConfig)
+    AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor)
+    AutoProcessor.register(OpenVLAConfig, PrismaticProcessor)
+    AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction)
+
+
+def load_vla_model(ckpt_path: str, device: str = DEFAULT_DEVICE) -> Tuple[AutoProcessor, AutoModelForVision2Seq]:
+    """
+    Load OpenVLA model and processor from checkpoint.
+    
+    Args:
+        ckpt_path: Path to the model checkpoint
+        device: Device to load the model on
+        
+    Returns:
+        Tuple of (processor, model)
+        
+    Raises:
+        FileNotFoundError: If checkpoint path or dataset statistics file doesn't exist
+    """
+    if not os.path.exists(ckpt_path):
+        raise FileNotFoundError(f"Checkpoint path does not exist: {ckpt_path}")
+    
+    stats_path = os.path.join(ckpt_path, "dataset_statistics.json")
+    if not os.path.exists(stats_path):
+        raise FileNotFoundError(f"Dataset statistics file not found: {stats_path}")
+    
+    processor = AutoProcessor.from_pretrained(ckpt_path, trust_remote_code=True)
+    model = AutoModelForVision2Seq.from_pretrained(
+        ckpt_path,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True
+    ).to(device)
+
+    with open(stats_path, "r") as f:
+        model.norm_stats = json.load(f)
+
+    return processor, model
+
+
+def setup_liftpot_env(downsample_rate: int = DEFAULT_DOWNSAMPLE_RATE) -> LiftPot:
+    """
+    Set up the LiftPot environment with bimanual robot configuration.
+    
+    Args:
+        downsample_rate: Control frequency downsampling factor
+        
+    Returns:
+        Configured LiftPot environment
+    """
+    return LiftPot(
+        action_mode=JointPositionActionMode(
+            floating_base=True, 
+            absolute=True, 
+            block_until_reached=False, 
+            ee=True, 
+            floating_dofs=[]
+        ),
+        observation_config=ObservationConfig(
+            cameras=[
+                CameraConfig(name="head", rgb=True, depth=False, resolution=CAMERA_RESOLUTION),
+                CameraConfig(name="left_wrist", rgb=True, depth=False, resolution=CAMERA_RESOLUTION),
+                CameraConfig(name="right_wrist", rgb=True, depth=False, resolution=CAMERA_RESOLUTION),
+            ]
+        ),
+        render_mode=None,
+        robot_cls=BimanualPanda,
+        control_frequency=CONTROL_FREQUENCY_MAX // downsample_rate,
+    )
+
+
+def get_successful_demos(
+    env: LiftPot, 
+    dataset_path: str, 
+    downsample_rate: int, 
+    num_demos: int = 1
+) -> List:
+    """
+    Load successful demonstrations from the dataset.
+    
+    Args:
+        env: Environment instance for metadata
+        dataset_path: Path to the demonstration dataset
+        downsample_rate: Frequency downsampling factor
+        num_demos: Number of demonstrations to load
+        
+    Returns:
+        List of successful demonstrations (reward > threshold)
+        
+    Raises:
+        FileNotFoundError: If dataset path doesn't exist
+    """
+    dataset_path = Path(dataset_path)
+    if not dataset_path.exists():
+        raise FileNotFoundError(f"Dataset path does not exist: {dataset_path}")
+    
+    demo_store = DemoStore()
+    demos = demo_store.get_demos_from_folder(
+        dataset_path,
+        Metadata.from_env(env),
+        amount=num_demos,
+        frequency=CONTROL_FREQUENCY_MAX // downsample_rate,
+    )
+    
+    successful_demos = [
+        copy.deepcopy(demo)
+        for demo in demos
+        if sum(step.reward for step in demo.timesteps) > MIN_REWARD_THRESHOLD
+    ]
+    
+    print(f"Found {len(successful_demos)} successful demos out of {len(demos)} total")
+    return successful_demos
+
+
+def run_inference_loop(
+    env: LiftPot,
+    processor: AutoProcessor,
+    model: AutoModelForVision2Seq,
+    instruction: str,
+    use_demos: bool = False,
+    demo: Optional[object] = None,
+    device: str = DEFAULT_DEVICE,
+    max_steps: int = DEFAULT_MAX_STEPS
+) -> List[np.ndarray]:
+    """
+    Run the agent in a closed loop using either demo replay or OpenVLA predictions.
+    
+    Args:
+        env: Environment instance
+        processor: Model processor for input preparation
+        model: OpenVLA model for action prediction
+        instruction: Task instruction string
+        use_demos: Whether to use demonstration actions instead of model predictions
+        demo: Demonstration object (required if use_demos=True)
+        device: Device for model inference
+        max_steps: Maximum number of steps to run
+        
+    Returns:
+        List of RGB frames from the episode
+    """
+    obs = env.reset(seed=demo.seed if (use_demos and demo) else None)
+    images = []
+    prompt = f"In: What action should the robot take to {instruction}?\nOut:"
+
+    episode_length = min(max_steps, len(demo.timesteps) if demo else max_steps)
+    
+    for step_idx in range(episode_length):
+        # Collect observation image
+        if step_idx == 0:
+            images.append(copy.deepcopy(obs[0]["rgb_head"]))
+        else:
+            images.append(obs["rgb_head"])
+
+        # Get action from demo or model
+        if use_demos and demo:
+            action = demo.timesteps[step_idx].executed_action
+        else:
+            image = Image.fromarray(np.moveaxis(images[-1], 0, -1))
+            inputs = processor(prompt, image).to(device, dtype=torch.bfloat16)
+            action = model.predict_action(**inputs, do_sample=False)
+
+        # Execute action
+        obs, reward, terminated, truncated, info = env.step(action)
+        
+        if terminated or truncated:
+            print(f"Episode ended at step {step_idx + 1} (terminated: {terminated}, truncated: {truncated})")
+            break
+
+    return images
+
+
+def visualize_frames(frames: List[np.ndarray], fps: int = DEFAULT_FPS) -> None:
+    """
+    Preview frames as a video using moviepy.
+    
+    Args:
+        frames: List of RGB frames (C, H, W format)
+        fps: Frames per second for playback
+    """
+    frames = np.moveaxis(np.array(frames), 1, -1)  # (T, C, H, W) → (T, H, W, C)
+    duration = len(frames) / fps
+    clip = VideoClip(make_frame=lambda t: frames[min(int(t * fps), len(frames) - 1)], duration=duration)
+    clip.preview()
+
+
+def save_frames_to_video(
+    frames: List[np.ndarray], 
+    output_dir: str, 
+    filename: str = "rollout.mp4", 
+    fps: int = DEFAULT_FPS
+) -> None:
+    """
+    Save rollout frames to a video file.
+    
+    Args:
+        frames: List of RGB frames (C, H, W format)
+        output_dir: Output directory path
+        filename: Output video filename
+        fps: Frames per second for the output video
+    """
+    output_path = os.path.join(output_dir, filename)
+    os.makedirs(output_dir, exist_ok=True)
+    
+    frames = np.moveaxis(np.array(frames), 1, -1)  # (T, C, H, W) → (T, H, W, C)
+    duration = len(frames) / fps
+    clip = VideoClip(make_frame=lambda t: frames[min(int(t * fps), len(frames) - 1)], duration=duration)
+    clip.write_videofile(output_path, fps=fps, codec="libx264")
+    print(f"Saved rollout video to: {output_path}")
+
+
+
+def main() -> None:
+    """Main evaluation function."""
+    # Parse arguments and create configuration
+    args = parse_arguments()
+    config = EvaluationConfig.from_args(args)
+    
+    # Validate configuration
+    try:
+        config.validate()
+    except ValueError as e:
+        print(f"Configuration error: {e}")
+        return
+    
+    # Set output directory
+    if config.output_dir is None:
+        config.output_dir = config.ckpt_path
+    
+    print("=== OpenVLA Evaluation Configuration ===")
+    for field_name in config.__dataclass_fields__:
+        value = getattr(config, field_name)
+        print(f"{field_name}: {value}")
+    print("=" * 40)
+
+    try:
+        # Register OpenVLA components
+        print("Registering OpenVLA components...")
+        register_openvla()
+
+        # Load model
+        print(f"Loading VLA model from: {config.ckpt_path}")
+        processor, model = load_vla_model(config.ckpt_path, config.device)
+
+        # Setup environment
+        print("Setting up environment...")
+        env = setup_liftpot_env(downsample_rate=config.downsample_rate)
+
+        # Load demonstrations if needed
+        demos = []
+        if config.use_demos:
+            print("Loading demonstration data...")
+            demos = get_successful_demos(
+                env, 
+                config.dataset_path, 
+                config.downsample_rate, 
+                num_demos=config.num_episodes
+            )
+            
+            if len(demos) < config.num_episodes:
+                print(f"Warning: Only found {len(demos)} successful demos, "
+                      f"reducing episodes to {len(demos)}")
+                config.num_episodes = len(demos)
+
+        # Run evaluation episodes
+        print(f"\nRunning {config.num_episodes} episodes...")
+        for ep in range(config.num_episodes):
+            mode = "demo" if config.use_demos else "inference"
+            print(f"\nEpisode {ep + 1}/{config.num_episodes} ({mode} mode)")
+            
+            demo = DemoConverter.joint_to_ee(demos[ep]) if config.use_demos else None
+            
+            # Run episode
+            frames = run_inference_loop(
+                env=env,
+                processor=processor,
+                model=model,
+                instruction=config.instruction,
+                use_demos=config.use_demos,
+                demo=demo,
+                device=config.device,
+                max_steps=config.max_steps
+            )
+            
+            # Save video
+            filename = f"rollout_{mode}_ep{ep + 1}.mp4"
+            save_frames_to_video(
+                frames, 
+                config.output_dir, 
+                filename=filename, 
+                fps=config.fps
+            )
+            print(f"Episode {ep + 1} completed with {len(frames)} frames")
+
+        print("\nEvaluation completed successfully!")
+        
+    except Exception as e:
+        print(f"Error during evaluation: {e}")
+        raise
+    finally:
+        # Cleanup
+        if 'env' in locals():
+            env.close()
+            print("Environment closed.")
+
+
+
+if __name__ == "__main__":
+    main()

inference_openpi.py

@@ -0,0 +1,431 @@
+"""
+OpenPI Inference Worker - Runs in openpi_env
+Receives inference requests via stdin, returns results via stdout
+"""
+import sys
+import json
+import os
+import tempfile
+import copy
+import numpy as np
+import dataclasses
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+
+# Set headless mode
+os.environ.setdefault("MUJOCO_GL", "egl")
+os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
+os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp")
+
+# Import OpenPI dependencies (only available in openpi_env)
+from openpi.training.config import Config as PIConfig
+from openpi.policies.policy_config import PIPolicy
+from roboeval.action_modes import JointPositionActionMode
+from roboeval.utils.observation_config import ObservationConfig, CameraConfig
+from roboeval.robots.configs.panda import BimanualPanda
+from roboeval.roboeval_env import CONTROL_FREQUENCY_MAX
+
+# Import all environment classes
+from roboeval.envs.manipulation import (
+    CubeHandover, CubeHandoverOrientation, CubeHandoverPosition,
+    CubeHandoverPositionAndOrientation, VerticalCubeHandover,
+    StackTwoBlocks, StackTwoBlocksOrientation, StackTwoBlocksPosition,
+    StackTwoBlocksPositionAndOrientation
+)
+from roboeval.envs.lift_pot import (
+    LiftPot, LiftPotOrientation, LiftPotPosition, LiftPotPositionAndOrientation,
+)
+from roboeval.envs.lift_tray import (
+    LiftTray, DragOverAndLiftTray, LiftTrayOrientation, LiftTrayPosition, LiftTrayPositionAndOrientation,
+)
+from roboeval.envs.pack_objects import (
+    PackBox, PackBoxOrientation, PackBoxPosition, PackBoxPositionAndOrientation,
+)
+from roboeval.envs.stack_books import (
+    PickSingleBookFromTable, PickSingleBookFromTableOrientation,
+    PickSingleBookFromTablePosition, PickSingleBookFromTablePositionAndOrientation,
+    StackSingleBookShelf, StackSingleBookShelfPosition, StackSingleBookShelfPositionAndOrientation,
+)
+from roboeval.envs.rotate_utility_objects import (
+    RotateValve, RotateValveObstacle, RotateValvePosition, RotateValvePositionAndOrientation,
+)
+
+# Video
+from moviepy.editor import VideoClip
+
+# Environment registry
+_ENV_CLASSES = {
+    "CubeHandover": (CubeHandover, "handover the rod from one hand to the other hand"),
+    "CubeHandoverOrientation": (CubeHandoverOrientation, "handover the rod from one hand to the other hand"),
+    "CubeHandoverPosition": (CubeHandoverPosition, "handover the rod from one hand to the other hand"),
+    "CubeHandoverPositionOrientation": (CubeHandoverPositionAndOrientation, "handover the rod from one hand to the other hand"),
+    "CubeHandoverVertical": (VerticalCubeHandover, "handover the rod from one hand to the other hand"),
+    "LiftPot": (LiftPot, "lift the pot by the handles"),
+    "LiftPotOrientation": (LiftPotOrientation, "lift the pot by the handles"),
+    "LiftPotPosition": (LiftPotPosition, "lift the pot by the handles"),
+    "LiftPotPositionOrientation": (LiftPotPositionAndOrientation, "lift the pot by the handles"),
+    "LiftTray": (LiftTray, "lift the tray"),
+    "LiftTrayDrag": (DragOverAndLiftTray, "lift the tray"),
+    "LiftTrayOrientation": (LiftTrayOrientation, "lift the tray"),
+    "LiftTrayPosition": (LiftTrayPosition, "lift the tray"),
+    "LiftTrayPositionOrientation": (LiftTrayPositionAndOrientation, "lift the tray"),
+    "PackBox": (PackBox, "close the box"),
+    "PackBoxOrientation": (PackBoxOrientation, "close the box"),
+    "PackBoxPosition": (PackBoxPosition, "close the box"),
+    "PackBoxPositionOrientation": (PackBoxPositionAndOrientation, "close the box"),
+    "PickSingleBookFromTable": (PickSingleBookFromTable, "pick up the book from the table"),
+    "PickSingleBookFromTableOrientation": (PickSingleBookFromTableOrientation, "pick up the book from the table"),
+    "PickSingleBookFromTablePosition": (PickSingleBookFromTablePosition, "pick up the book from the table"),
+    "PickSingleBookFromTablePositionOrientation": (PickSingleBookFromTablePositionAndOrientation, "pick up the book from the table"),
+    "RotateValve": (RotateValve, "rotate the valve counter clockwise"),
+    "RotateValveObstacle": (RotateValveObstacle, "rotate the valve counter clockwise"),
+    "RotateValvePosition": (RotateValvePosition, "rotate the valve counter clockwise"),
+    "RotateValvePositionOrientation": (RotateValvePositionAndOrientation, "rotate the valve counter clockwise"),
+    "StackSingleBookShelf": (StackSingleBookShelf, "put the book on the table onto the shelf"),
+    "StackSingleBookShelfPosition": (StackSingleBookShelfPosition, "put the book on the table onto the shelf"),
+    "StackSingleBookShelfPositionOrientation": (StackSingleBookShelfPositionAndOrientation, "put the book on the table onto the shelf"),
+    "StackTwoBlocks": (StackTwoBlocks, "stack the two cubes"),
+    "StackTwoBlocksOrientation": (StackTwoBlocksOrientation, "stack the two cubes"),
+    "StackTwoBlocksPosition": (StackTwoBlocksPosition, "stack the two cubes"),
+    "StackTwoBlocksPositionOrientation": (StackTwoBlocksPositionAndOrientation, "stack the two cubes")
+}
+
+DEFAULT_DOWNSAMPLE_RATE = 25
+
+# Global policy cache
+_POLICY_CACHE = {}
+
+
+def get_checkpoint_path(task_name: str, ckpt_path: Optional[str] = None) -> str:
+    """
+    Return a local path to the checkpoint for the given task.
+    """
+    if ckpt_path:
+        return ckpt_path
+
+    from huggingface_hub import HfApi, hf_hub_download
+
+    repo_id = "tan7271/pi0_base_checkpoints"
+    revision = "main"
+    base_dir = f"{task_name}_testing"
+    cache_dir = os.path.expanduser("~/.cache/roboeval/pi0_checkpoints")
+
+    api = HfApi()
+    try:
+        all_files: List[str] = api.list_repo_files(
+            repo_id=repo_id, revision=revision, repo_type="model"
+        )
+    except Exception as e:
+        raise RuntimeError(f"Could not list files for {repo_id}@{revision}: {e}")
+
+    steps = sorted({
+        int(p.split("/")[1])
+        for p in all_files
+        if p.startswith(base_dir + "/") and len(p.split("/")) >= 3 and p.split("/")[1].isdigit()
+    })
+    if not steps:
+        nearby = [p for p in all_files if base_dir in p][:10]
+        raise FileNotFoundError(
+            f"No files found under '{base_dir}/' in {repo_id}@{revision}. "
+            f"Example paths I do see: {nearby}"
+        )
+
+    chosen_step = 2999 if 2999 in steps else steps[-1]
+    subdir = f"{base_dir}/{chosen_step}"
+
+    needed_roots = (
+        f"{subdir}/_CHECKPOINT_METADATA",
+        f"{subdir}/assets/",
+        f"{subdir}/params/",
+        f"{subdir}/train_state/",
+    )
+    wanted = [
+        p for p in all_files
+        if p == f"{subdir}/_CHECKPOINT_METADATA"
+        or any(p.startswith(root) for root in needed_roots[1:])
+    ]
+
+    if not wanted:
+        wanted = [p for p in all_files if p.startswith(subdir + "/")]
+        if not wanted:
+            raise FileNotFoundError(
+                f"Repo listing shows no files under '{subdir}/'. "
+                f"Steps seen: {steps}"
+            )
+
+    manual_root = os.path.join(cache_dir, "manual")
+    os.makedirs(manual_root, exist_ok=True)
+
+    for relpath in wanted:
+        hf_hub_download(
+            repo_id=repo_id,
+            filename=relpath,
+            revision=revision,
+            repo_type="model",
+            local_dir=manual_root,
+            local_dir_use_symlinks=True,
+        )
+
+    manual_ckpt_dir = os.path.join(manual_root, subdir)
+
+    def _nonempty_dir(path: str) -> bool:
+        return os.path.isdir(path) and any(True for _ in os.scandir(path))
+
+    if not _nonempty_dir(manual_ckpt_dir):
+        try:
+            siblings = [e.name for e in os.scandir(os.path.dirname(manual_ckpt_dir))]
+        except Exception:
+            siblings = []
+        raise FileNotFoundError(
+            f"Downloaded files, but '{manual_ckpt_dir}' is missing/empty.\n"
+            f"Siblings present: {siblings}\n"
+            f"(repo_id={repo_id}, subdir={subdir})"
+        )
+
+    return manual_ckpt_dir
+
+
+def load_pi0_policy(task_name: str, ckpt_path: str):
+    """Load Pi0 policy model for the given task."""
+    checkpoint_path = get_checkpoint_path(task_name, ckpt_path)
+    
+    cache_key = f"{task_name}:{checkpoint_path}"
+    if cache_key in _POLICY_CACHE:
+        return _POLICY_CACHE[cache_key]
+    
+    cfg = PIConfig.get_config("pi0_base_bimanual_droid_finetune")
+    bimanual_assets = PIConfig.AssetsConfig(
+        assets_dir=f"{checkpoint_path}/assets/",
+        asset_id=f"tan7271/{task_name}",
+    )
+    cfg = dataclasses.replace(cfg, data=dataclasses.replace(cfg.data, assets=bimanual_assets))
+    policy = PIPolicy.create_trained_policy(cfg, checkpoint_path)
+    
+    _POLICY_CACHE[cache_key] = policy
+    return policy
+
+
+def make_openpi_example_from_roboeval(obs_dict: dict, prompt: str) -> dict:
+    """Convert RoboEval observation to OpenPI format."""
+    obs = obs_dict[0] if isinstance(obs_dict, (tuple, list)) else obs_dict
+    example = {"prompt": prompt}
+
+    exterior_chw = obs["rgb_head"]
+    left_wrist_chw = obs["rgb_left_wrist"]
+    right_wrist_chw = obs["rgb_right_wrist"]
+    example["observation/exterior_image_1_left"] = np.moveaxis(exterior_chw, 0, -1)
+    example["observation/wrist_image_left"] = np.moveaxis(left_wrist_chw, 0, -1)
+    example["observation/wrist_image_right"] = np.moveaxis(right_wrist_chw, 0, -1)
+
+    prop = np.asarray(obs["proprioception"], dtype=np.float32).reshape(-1)
+    example["observation/joint_position"] = prop
+
+    grip = np.asarray(obs["proprioception_grippers"], dtype=np.float32).reshape(-1)[:2]
+    example["observation/gripper_position"] = grip
+
+    return example
+
+
+def map_policy_action_to_env_abs(action_vec: np.ndarray, env) -> np.ndarray:
+    """Map policy action to environment action."""
+    a = np.asarray(action_vec, dtype=np.float32).reshape(-1)
+    if a.shape[0] != 16:
+        raise ValueError(f"Expected (16,), got {a.shape}.")
+    return a
+
+
+def _clip_to_space(env, action: np.ndarray) -> np.ndarray:
+    """Safety: clip to env action space."""
+    return np.clip(action, env.action_space.low, env.action_space.high)
+
+
+def setup_env(task_name: str, downsample_rate: int = 25):
+    """Setup RoboEval environment for the given task."""
+    cameras = [
+        CameraConfig(name="head", rgb=True, depth=False, resolution=(256, 256)),
+        CameraConfig(name="left_wrist", rgb=True, depth=False, resolution=(256, 256)),
+        CameraConfig(name="right_wrist", rgb=True, depth=False, resolution=(256, 256)),
+    ]
+    
+    Env, prompt = _ENV_CLASSES[task_name]
+    env = Env(
+        action_mode=JointPositionActionMode(
+            floating_base=True,
+            absolute=False,
+            block_until_reached=False,
+            ee=False,
+            floating_dofs=[],
+        ),
+        observation_config=ObservationConfig(cameras=cameras, proprioception=True),
+        render_mode="rgb_array",
+        robot_cls=BimanualPanda,
+        control_frequency=CONTROL_FREQUENCY_MAX // downsample_rate,
+    )
+    
+    return env, prompt
+
+
+def _unpack_obs(obs_or_tuple):
+    """Unpack observation if it's a tuple."""
+    return obs_or_tuple[0] if isinstance(obs_or_tuple, (tuple, list)) else obs_or_tuple
+
+
+def run_inference_loop(
+    env,
+    policy,
+    instruction: str,
+    max_steps: int = 200,
+    open_loop_horizon: int = 10,
+):
+    """Run inference loop for one episode."""
+    obs = env.reset()
+    successes = 0
+    images_env = []
+    chunk = None
+    i_in_chunk = 0
+
+    for step_idx in range(max_steps):
+        cur_obs = _unpack_obs(obs)
+        
+        images_env.append(copy.deepcopy(env.render()))
+        
+        if chunk is None or i_in_chunk >= open_loop_horizon:
+            example = make_openpi_example_from_roboeval(cur_obs, instruction)
+            out = policy.infer(example)
+            chunk = out["actions"]
+            i_in_chunk = 0
+
+        a_vec = chunk[i_in_chunk]
+        i_in_chunk += 1
+
+        env_action = map_policy_action_to_env_abs(a_vec, env)
+        env_action = _clip_to_space(env, env_action)
+
+        obs, reward, terminated, truncated, info = env.step(env_action)
+        successes += int(reward > 0)
+
+        if terminated or truncated:
+            break
+
+    stats = {"steps": step_idx + 1, "success_signal": successes}
+    return stats, images_env
+
+
+def save_frames_to_video(frames, output_path: str, fps: int = 25) -> str:
+    """Save rollout frames to a video file."""
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    
+    frames = np.array(frames)
+    if frames.ndim == 4 and frames.shape[1] == 3:
+        frames = np.moveaxis(frames, 1, -1)
+    
+    duration = len(frames) / fps
+    clip = VideoClip(make_frame=lambda t: frames[min(int(t * fps), len(frames) - 1)], duration=duration)
+    clip.write_videofile(output_path, fps=fps, codec="libx264", logger=None)
+    
+    return output_path
+
+
+def run_inference(request: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Run OpenPI inference based on request parameters.
+    
+    Args:
+        request: Dictionary with keys:
+            - task_name: str
+            - checkpoint_path: str or None
+            - max_steps: int
+            - fps: int
+            - custom_instruction: str or None
+    
+    Returns:
+        Dictionary with keys:
+            - success: bool
+            - video_path: str or None
+            - status_message: str
+            - error: str or None
+    """
+    try:
+        task_name = request["task_name"]
+        checkpoint_path = request.get("checkpoint_path") or None
+        max_steps = int(request["max_steps"])
+        fps = int(request["fps"])
+        custom_instruction = request.get("custom_instruction") or None
+        
+        # Validate task
+        if task_name not in _ENV_CLASSES:
+            return {
+                "success": False,
+                "video_path": None,
+                "status_message": f"❌ Unknown task: {task_name}",
+                "error": f"Unknown task: {task_name}"
+            }
+        
+        # Load policy
+        policy = load_pi0_policy(task_name, checkpoint_path or "")
+        
+        # Setup environment
+        env, default_prompt = setup_env(task_name, downsample_rate=DEFAULT_DOWNSAMPLE_RATE)
+        instruction = custom_instruction if custom_instruction else default_prompt
+        
+        # Run inference
+        stats, images_env = run_inference_loop(
+            env, policy, instruction, max_steps=max_steps
+        )
+        
+        # Save video
+        video_path = os.path.join(tempfile.gettempdir(), f"pi0_rollout_{task_name}_{os.getpid()}.mp4")
+        save_frames_to_video(images_env, video_path, fps=fps)
+        
+        # Cleanup
+        env.close()
+        
+        status = f"✅ **Inference Complete!**\n\n"
+        status += f"- **Task**: {task_name}\n"
+        status += f"- **Steps**: {stats['steps']}\n"
+        status += f"- **Success Signal**: {stats['success_signal']}\n"
+        status += f"- **Instruction**: {instruction}\n"
+        
+        return {
+            "success": True,
+            "video_path": video_path,
+            "status_message": status,
+            "error": None
+        }
+        
+    except Exception as e:
+        import traceback
+        error_msg = f"❌ **Error during inference:**\n\n```\n{str(e)}\n\n{traceback.format_exc()}\n```"
+        return {
+            "success": False,
+            "video_path": None,
+            "status_message": error_msg,
+            "error": str(e)
+        }
+
+
+def main():
+    """Main loop: read requests from stdin, write results to stdout"""
+    while True:
+        try:
+            line = sys.stdin.readline()
+            if not line:
+                break
+            
+            request = json.loads(line.strip())
+            result = run_inference(request)
+            print(json.dumps(result), flush=True)
+            
+        except Exception as e:
+            error_result = {
+                "success": False,
+                "video_path": None,
+                "status_message": "❌ Worker error",
+                "error": str(e)
+            }
+            print(json.dumps(error_result), flush=True)
+
+
+if __name__ == "__main__":
+    main()
+

inference_openvla.py

@@ -0,0 +1,351 @@
+"""
+OpenVLA Inference Worker - Runs in openvla_env
+Receives inference requests via stdin, returns results via stdout
+"""
+import sys
+import json
+import os
+import tempfile
+import copy
+import numpy as np
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+
+# Set headless mode
+os.environ.setdefault("MUJOCO_GL", "egl")
+os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
+os.environ.setdefault("XDG_RUNTIME_DIR", "/tmp")
+
+import torch
+from PIL import Image
+from transformers import (
+    AutoConfig,
+    AutoImageProcessor,
+    AutoModelForVision2Seq,
+    AutoProcessor,
+)
+
+# Import OpenVLA dependencies (only available in openvla_env)
+from prismatic.extern.hf.configuration_prismatic import OpenVLAConfig
+from prismatic.extern.hf.modeling_prismatic import OpenVLAForActionPrediction
+from prismatic.extern.hf.processing_prismatic import PrismaticImageProcessor, PrismaticProcessor
+
+from roboeval.action_modes import JointPositionActionMode
+from roboeval.utils.observation_config import CameraConfig, ObservationConfig
+from roboeval.robots.configs.panda import BimanualPanda
+from roboeval.roboeval_env import CONTROL_FREQUENCY_MAX
+
+# Import all environment classes
+from roboeval.envs.manipulation import (
+    CubeHandover, CubeHandoverOrientation, CubeHandoverPosition,
+    CubeHandoverPositionAndOrientation, VerticalCubeHandover,
+    StackTwoBlocks, StackTwoBlocksOrientation, StackTwoBlocksPosition,
+    StackTwoBlocksPositionAndOrientation
+)
+from roboeval.envs.lift_pot import (
+    LiftPot, LiftPotOrientation, LiftPotPosition, LiftPotPositionAndOrientation,
+)
+from roboeval.envs.lift_tray import (
+    LiftTray, DragOverAndLiftTray, LiftTrayOrientation, LiftTrayPosition, LiftTrayPositionAndOrientation,
+)
+from roboeval.envs.pack_objects import (
+    PackBox, PackBoxOrientation, PackBoxPosition, PackBoxPositionAndOrientation,
+)
+from roboeval.envs.stack_books import (
+    PickSingleBookFromTable, PickSingleBookFromTableOrientation,
+    PickSingleBookFromTablePosition, PickSingleBookFromTablePositionAndOrientation,
+    StackSingleBookShelf, StackSingleBookShelfPosition, StackSingleBookShelfPositionAndOrientation,
+)
+from roboeval.envs.rotate_utility_objects import (
+    RotateValve, RotateValveObstacle, RotateValvePosition, RotateValvePositionAndOrientation,
+)
+
+from moviepy.editor import VideoClip
+
+# Configuration constants
+DEFAULT_DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
+DEFAULT_DOWNSAMPLE_RATE = 25
+CAMERA_RESOLUTION = (256, 256)
+
+# Environment registry
+_ENV_CLASSES = {
+    "CubeHandover": (CubeHandover, "handover the rod from one hand to the other hand"),
+    "CubeHandoverOrientation": (CubeHandoverOrientation, "handover the rod from one hand to the other hand"),
+    "CubeHandoverPosition": (CubeHandoverPosition, "handover the rod from one hand to the other hand"),
+    "CubeHandoverPositionOrientation": (CubeHandoverPositionAndOrientation, "handover the rod from one hand to the other hand"),
+    "CubeHandoverVertical": (VerticalCubeHandover, "handover the rod from one hand to the other hand"),
+    "LiftPot": (LiftPot, "lift the pot by the handles"),
+    "LiftPotOrientation": (LiftPotOrientation, "lift the pot by the handles"),
+    "LiftPotPosition": (LiftPotPosition, "lift the pot by the handles"),
+    "LiftPotPositionOrientation": (LiftPotPositionAndOrientation, "lift the pot by the handles"),
+    "LiftTray": (LiftTray, "lift the tray"),
+    "LiftTrayDrag": (DragOverAndLiftTray, "lift the tray"),
+    "LiftTrayOrientation": (LiftTrayOrientation, "lift the tray"),
+    "LiftTrayPosition": (LiftTrayPosition, "lift the tray"),
+    "LiftTrayPositionOrientation": (LiftTrayPositionAndOrientation, "lift the tray"),
+    "PackBox": (PackBox, "close the box"),
+    "PackBoxOrientation": (PackBoxOrientation, "close the box"),
+    "PackBoxPosition": (PackBoxPosition, "close the box"),
+    "PackBoxPositionOrientation": (PackBoxPositionAndOrientation, "close the box"),
+    "PickSingleBookFromTable": (PickSingleBookFromTable, "pick up the book from the table"),
+    "PickSingleBookFromTableOrientation": (PickSingleBookFromTableOrientation, "pick up the book from the table"),
+    "PickSingleBookFromTablePosition": (PickSingleBookFromTablePosition, "pick up the book from the table"),
+    "PickSingleBookFromTablePositionOrientation": (PickSingleBookFromTablePositionAndOrientation, "pick up the book from the table"),
+    "RotateValve": (RotateValve, "rotate the valve counter clockwise"),
+    "RotateValveObstacle": (RotateValveObstacle, "rotate the valve counter clockwise"),
+    "RotateValvePosition": (RotateValvePosition, "rotate the valve counter clockwise"),
+    "RotateValvePositionOrientation": (RotateValvePositionAndOrientation, "rotate the valve counter clockwise"),
+    "StackSingleBookShelf": (StackSingleBookShelf, "put the book on the table onto the shelf"),
+    "StackSingleBookShelfPosition": (StackSingleBookShelfPosition, "put the book on the table onto the shelf"),
+    "StackSingleBookShelfPositionOrientation": (StackSingleBookShelfPositionAndOrientation, "put the book on the table onto the shelf"),
+    "StackTwoBlocks": (StackTwoBlocks, "stack the two cubes"),
+    "StackTwoBlocksOrientation": (StackTwoBlocksOrientation, "stack the two cubes"),
+    "StackTwoBlocksPosition": (StackTwoBlocksPosition, "stack the two cubes"),
+    "StackTwoBlocksPositionOrientation": (StackTwoBlocksPositionAndOrientation, "stack the two cubes")
+}
+
+# Global model cache
+_MODEL_CACHE = {}
+
+
+def register_openvla() -> None:
+    """Register OpenVLA components with the Transformers library."""
+    AutoConfig.register("openvla", OpenVLAConfig)
+    AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor)
+    AutoProcessor.register(OpenVLAConfig, PrismaticProcessor)
+    AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction)
+
+
+def load_vla_model(ckpt_path: str, device: str = DEFAULT_DEVICE) -> Tuple[AutoProcessor, AutoModelForVision2Seq]:
+    """
+    Load OpenVLA model and processor from checkpoint.
+    """
+    if ckpt_path in _MODEL_CACHE:
+        return _MODEL_CACHE[ckpt_path]
+    
+    if not os.path.exists(ckpt_path):
+        raise FileNotFoundError(f"Checkpoint path does not exist: {ckpt_path}")
+    
+    stats_path = os.path.join(ckpt_path, "dataset_statistics.json")
+    if not os.path.exists(stats_path):
+        raise FileNotFoundError(f"Dataset statistics file not found: {stats_path}")
+    
+    processor = AutoProcessor.from_pretrained(ckpt_path, trust_remote_code=True)
+    model = AutoModelForVision2Seq.from_pretrained(
+        ckpt_path,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True
+    ).to(device)
+
+    with open(stats_path, "r") as f:
+        model.norm_stats = json.load(f)
+
+    _MODEL_CACHE[ckpt_path] = (processor, model)
+    return processor, model
+
+
+def setup_env(task_name: str, downsample_rate: int = DEFAULT_DOWNSAMPLE_RATE):
+    """Setup RoboEval environment for the given task."""
+    cameras = [
+        CameraConfig(name="head", rgb=True, depth=False, resolution=CAMERA_RESOLUTION),
+        CameraConfig(name="left_wrist", rgb=True, depth=False, resolution=CAMERA_RESOLUTION),
+        CameraConfig(name="right_wrist", rgb=True, depth=False, resolution=CAMERA_RESOLUTION),
+    ]
+    
+    Env, prompt = _ENV_CLASSES[task_name]
+    env = Env(
+        action_mode=JointPositionActionMode(
+            floating_base=True,
+            absolute=True,
+            block_until_reached=False,
+            ee=True,
+            floating_dofs=[],
+        ),
+        observation_config=ObservationConfig(cameras=cameras, proprioception=True),
+        render_mode="rgb_array",
+        robot_cls=BimanualPanda,
+        control_frequency=CONTROL_FREQUENCY_MAX // downsample_rate,
+    )
+    
+    return env, prompt
+
+
+def _unpack_obs(obs_or_tuple):
+    """Unpack observation if it's a tuple."""
+    return obs_or_tuple[0] if isinstance(obs_or_tuple, (tuple, list)) else obs_or_tuple
+
+
+def run_inference_loop(
+    env,
+    processor: AutoProcessor,
+    model: AutoModelForVision2Seq,
+    instruction: str,
+    device: str = DEFAULT_DEVICE,
+    max_steps: int = 200
+) -> Tuple[Dict[str, Any], List[np.ndarray]]:
+    """
+    Run the agent in a closed loop using OpenVLA predictions.
+    
+    Returns:
+        Tuple of (stats dict, list of RGB frames)
+    """
+    obs = env.reset()
+    images = []
+    prompt = f"In: What action should the robot take to {instruction}?\nOut:"
+    successes = 0
+
+    for step_idx in range(max_steps):
+        # Collect observation image
+        cur_obs = _unpack_obs(obs)
+        images.append(copy.deepcopy(cur_obs["rgb_head"]))
+        
+        # Get action from model
+        image = Image.fromarray(np.moveaxis(images[-1], 0, -1))
+        inputs = processor(prompt, image).to(device, dtype=torch.bfloat16)
+        action = model.predict_action(**inputs, do_sample=False)
+
+        # Execute action
+        obs, reward, terminated, truncated, info = env.step(action)
+        successes += int(reward > 0)
+        
+        if terminated or truncated:
+            break
+
+    stats = {"steps": step_idx + 1, "success_signal": successes}
+    return stats, images
+
+
+def save_frames_to_video(frames: List[np.ndarray], output_path: str, fps: int = 25) -> str:
+    """Save rollout frames to a video file."""
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    
+    frames = np.moveaxis(np.array(frames), 1, -1)  # (T, C, H, W) → (T, H, W, C)
+    duration = len(frames) / fps
+    clip = VideoClip(make_frame=lambda t: frames[min(int(t * fps), len(frames) - 1)], duration=duration)
+    clip.write_videofile(output_path, fps=fps, codec="libx264", logger=None)
+    
+    return output_path
+
+
+def run_inference(request: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Run OpenVLA inference based on request parameters.
+    
+    Args:
+        request: Dictionary with keys:
+            - task_name: str
+            - checkpoint_path: str (required)
+            - max_steps: int
+            - fps: int
+            - custom_instruction: str or None
+    
+    Returns:
+        Dictionary with keys:
+            - success: bool
+            - video_path: str or None
+            - status_message: str
+            - error: str or None
+    """
+    try:
+        task_name = request["task_name"]
+        checkpoint_path = request.get("checkpoint_path")
+        max_steps = int(request["max_steps"])
+        fps = int(request["fps"])
+        custom_instruction = request.get("custom_instruction") or None
+        
+        # Validate checkpoint path
+        if not checkpoint_path:
+            return {
+                "success": False,
+                "video_path": None,
+                "status_message": "❌ Checkpoint path is required for OpenVLA",
+                "error": "Checkpoint path is required"
+            }
+        
+        # Validate task
+        if task_name not in _ENV_CLASSES:
+            return {
+                "success": False,
+                "video_path": None,
+                "status_message": f"❌ Unknown task: {task_name}",
+                "error": f"Unknown task: {task_name}"
+            }
+        
+        # Register OpenVLA components
+        register_openvla()
+        
+        # Load model
+        device = DEFAULT_DEVICE
+        processor, model = load_vla_model(checkpoint_path, device)
+        
+        # Setup environment
+        env, default_prompt = setup_env(task_name, downsample_rate=DEFAULT_DOWNSAMPLE_RATE)
+        instruction = custom_instruction if custom_instruction else default_prompt
+        
+        # Run inference
+        stats, images = run_inference_loop(
+            env=env,
+            processor=processor,
+            model=model,
+            instruction=instruction,
+            device=device,
+            max_steps=max_steps
+        )
+        
+        # Save video
+        video_path = os.path.join(tempfile.gettempdir(), f"openvla_rollout_{task_name}_{os.getpid()}.mp4")
+        save_frames_to_video(images, video_path, fps=fps)
+        
+        # Cleanup
+        env.close()
+        
+        status = f"✅ **OpenVLA Inference Complete!**\n\n"
+        status += f"- **Task**: {task_name}\n"
+        status += f"- **Steps**: {stats['steps']}\n"
+        status += f"- **Success Signal**: {stats['success_signal']}\n"
+        status += f"- **Instruction**: {instruction}\n"
+        
+        return {
+            "success": True,
+            "video_path": video_path,
+            "status_message": status,
+            "error": None
+        }
+        
+    except Exception as e:
+        import traceback
+        error_msg = f"❌ **Error during OpenVLA inference:**\n\n```\n{str(e)}\n\n{traceback.format_exc()}\n```"
+        return {
+            "success": False,
+            "video_path": None,
+            "status_message": error_msg,
+            "error": str(e)
+        }
+
+
+def main():
+    """Main loop: read requests from stdin, write results to stdout"""
+    while True:
+        try:
+            line = sys.stdin.readline()
+            if not line:
+                break
+            
+            request = json.loads(line.strip())
+            result = run_inference(request)
+            print(json.dumps(result), flush=True)
+            
+        except Exception as e:
+            error_result = {
+                "success": False,
+                "video_path": None,
+                "status_message": "❌ Worker error",
+                "error": str(e)
+            }
+            print(json.dumps(error_result), flush=True)
+
+
+if __name__ == "__main__":
+    main()
+

requirements.txt

@@ -1,88 +1,12 @@
-# Core frameworks
-gradio>=4.0.0
+# Base environment (Gradio only)
+# Minimal dependencies for the main Gradio app
+# Model-specific dependencies are installed in separate conda environments
 
-# Essential OpenPI dependencies (with specific versions to avoid conflicts)
-augmax>=0.3.4
-dm-tree>=0.1.8
-einops>=0.8.0
-equinox>=0.11.8
-flatbuffers>=24.3.25
-flax==0.10.2
-fsspec[gcs]>=2024.6.0
-gym-aloha>=0.1.1
-imageio>=2.36.1
-jax[cuda12]==0.5.3
-jaxtyping==0.2.36
-ml_collections==1.0.0
+gradio>=4.0.0
 numpy>=1.22.4,<2.0.0
-numpydantic>=1.6.6
-opencv-python>=4.10.0.84
-orbax-checkpoint==0.11.13
 pillow>=11.0.0
-sentencepiece>=0.2.0
-torch>=2.4.0
-torchvision>=0.19.0
-torchaudio>=2.4.0
-tqdm-loggable>=0.2
-typing-extensions>=4.12.2
-tyro>=0.9.5
-wandb>=0.19.1
-filelock>=3.16.1
-beartype==0.19.0
-treescope>=0.1.7
-transformers==4.48.1
-rich>=14.0.0
-polars>=1.30.0
-
-# JAX and ML utilities
-jaxlib==0.5.3
-optax==0.2.5
-chex==0.1.90
-
-# Video/media processing
-moviepy==1.0.3
-imageio-ffmpeg==0.6.0
-opencv-python-headless==4.11.0.86
-
-# MuJoCo and robotics
-mujoco==3.3.3
-dm-control==1.0.31
-mujoco-utils==0.0.6
-gymnasium==0.29.1
-
-# Hugging Face (simplified versions)
-datasets==3.6.0
 huggingface-hub>=0.36.0
-diffusers==0.26.3
-tokenizers>=0.19.1
-
-# Utilities
-hydra-core==1.3.2
-omegaconf==2.3.0
-pyyaml==6.0.2
-pyquaternion==0.9.9
-click==8.2.1
-click-prompt==0.6.5
-rich-click==1.8.9
-
-# Additional dependencies
-h5py==3.14.0
-scipy==1.16.1
-pandas==2.3.2
-tqdm==4.67.1
-requests==2.32.5
-packaging==25.0
-safetensors>=0.4.1
-
-# openpi-client dependencies (needed since we install OpenPI with --no-deps)
-msgpack>=1.0.5
-tree>=0.2.4
-websockets>=11.0
-
-# OpenVLA dependencies (core packages needed for OpenVLA)
-accelerate>=1.11.0
-peft>=0.11.1
-timm>=0.9.10
 
-# Install lerobot and openvla separately to avoid conflicts
-# lerobot and openvla will be installed via git in setup.sh
+# Note: OpenPI and OpenVLA dependencies are installed in separate conda environments
+# See requirements_openpi.txt and requirements_openvla.txt
+# Setup is handled by setup.sh which creates openpi_env and openvla_env

requirements_openpi.txt

@@ -0,0 +1,81 @@
+# OpenPI environment dependencies
+# PyTorch 2.7.0 for OpenPI
+torch==2.7.0
+torchvision==0.20.0
+torchaudio==2.7.0
+
+# JAX and ML utilities
+jax[cuda12]==0.5.3
+jaxlib==0.5.3
+flax==0.10.2
+optax==0.2.5
+chex==0.1.90
+
+# Essential OpenPI dependencies
+augmax>=0.3.4
+dm-tree>=0.1.8
+einops>=0.8.0
+equinox>=0.11.8
+flatbuffers>=24.3.25
+fsspec[gcs]>=2024.6.0
+gym-aloha>=0.1.1
+imageio>=2.36.1
+jaxtyping==0.2.36
+ml_collections==1.0.0
+numpy>=1.22.4,<2.0.0
+numpydantic>=1.6.6
+opencv-python>=4.10.0.84
+orbax-checkpoint==0.11.13
+pillow>=11.0.0
+sentencepiece>=0.2.0
+tqdm-loggable>=0.2
+typing-extensions>=4.12.2
+tyro>=0.9.5
+wandb>=0.19.1
+filelock>=3.16.1
+beartype==0.19.0
+treescope>=0.1.7
+transformers==4.48.1
+rich>=14.0.0
+polars>=1.30.0
+
+# Hugging Face
+datasets==3.6.0
+huggingface-hub>=0.36.0
+diffusers==0.26.3
+tokenizers>=0.19.1
+
+# Utilities
+hydra-core==1.3.2
+omegaconf==2.3.0
+pyyaml==6.0.2
+pyquaternion==0.9.9
+click==8.2.1
+click-prompt==0.6.5
+rich-click==1.8.9
+
+# Additional dependencies
+h5py==3.14.0
+scipy==1.16.1
+pandas==2.3.2
+tqdm==4.67.1
+requests==2.32.5
+packaging==25.0
+safetensors>=0.4.1
+
+# openpi-client dependencies (needed since we install OpenPI with --no-deps)
+msgpack>=1.0.5
+tree>=0.2.4
+websockets>=11.0
+
+# Video/media processing
+moviepy==1.0.3
+imageio-ffmpeg==0.6.0
+opencv-python-headless==4.11.0.86
+
+# MuJoCo and robotics
+mujoco==3.3.3
+dm-control==1.0.31
+mujoco-utils==0.0.6
+gymnasium==0.29.1
+

requirements_openvla.txt

@@ -0,0 +1,42 @@
+# OpenVLA environment dependencies
+# PyTorch 2.2.0 for OpenVLA (installed separately in setup.sh with CUDA 12.1 index)
+# torch==2.2.0
+# torchvision==0.17.0
+# torchaudio==2.2.0
+
+# Transformers and OpenVLA dependencies
+transformers==4.40.1
+accelerate>=1.11.0
+peft>=0.11.1
+timm>=0.9.10
+
+# Hugging Face
+huggingface-hub>=0.36.0
+tokenizers>=0.19.1
+datasets==3.6.0
+
+# Core utilities
+numpy>=1.22.4,<2.0.0
+pillow>=11.0.0
+tqdm==4.67.1
+requests==2.32.5
+packaging==25.0
+
+# Video/media processing
+moviepy==1.0.3
+imageio-ffmpeg==0.6.0
+opencv-python-headless==4.11.0.86
+
+# MuJoCo and robotics
+mujoco==3.3.3
+dm-control==1.0.31
+mujoco-utils==0.0.6
+gymnasium==0.29.1
+
+# Additional utilities
+pyyaml==6.0.2
+pyquaternion==0.9.9
+scipy==1.16.1
+pandas==2.3.2
+h5py==3.14.0
+

setup.sh

@@ -1,10 +1,12 @@
 #!/bin/bash
 set -e
 
-echo "===== Installing Dependencies ====="
+echo "===== Multi-Environment Setup ====="
+echo "Building OpenPI environment (OpenVLA temporarily disabled)..."
 
-# Install RoboEval with submodules
-echo "Installing RoboEval with submodules..."
+# Install RoboEval in base (shared by both)
+echo ""
+echo "===== Installing RoboEval (shared) ====="
 CLONE_DIR="/tmp/roboeval_install"
 rm -rf $CLONE_DIR
 
@@ -16,40 +18,66 @@ git clone --recurse-submodules https://${GH_TOKEN}@github.com/helen9975/RoboEval
 echo "Installing RoboEval from cloned repository..."
 pip install $CLONE_DIR --no-cache-dir
 
-# Copy thirdparty to site-packages
-echo "Copying thirdparty submodules to site-packages..."
 SITE_PACKAGES=$(python -c "import site; print(site.getsitepackages()[0])")
 cp -r $CLONE_DIR/thirdparty $SITE_PACKAGES/
-echo "Copied thirdparty to $SITE_PACKAGES/thirdparty"
+echo "RoboEval installed in base environment"
 
-echo "RoboEval installed successfully with submodules"
+# Create OpenPI environment
+echo ""
+echo "===== Creating OpenPI Environment ====="
+conda create -n openpi_env python=3.10 -y
+conda run -n openpi_env pip install -r requirements_openpi.txt --no-cache-dir
 
-# Install lerobot from specific commit
-echo "Installing lerobot from git (specific commit required by OpenPI)..."
-pip install git+https://github.com/huggingface/lerobot@0cf864870cf29f4738d3ade893e6fd13fbd7cdb5 --no-cache-dir
+# Install lerobot and OpenPI in openpi_env
+echo "Installing lerobot in openpi_env..."
+conda run -n openpi_env pip install git+https://github.com/huggingface/lerobot@0cf864870cf29f4738d3ade893e6fd13fbd7cdb5 --no-cache-dir
 echo "lerobot installed successfully"
 
-# Upgrade safetensors to fix version conflict
-echo "Upgrading safetensors to >=0.4.1..."
-pip install "safetensors>=0.4.1" --upgrade --no-cache-dir
+echo "Upgrading safetensors in openpi_env..."
+conda run -n openpi_env pip install "safetensors>=0.4.1" --upgrade --no-cache-dir
 echo "safetensors upgraded successfully"
 
-# Install OpenPI with openpi-client
-echo "Installing OpenPI from tan7271/OpenPiRoboEval..."
-
-# Install openpi-client
-echo "Installing openpi-client..."
-pip install git+https://${GH_TOKEN}@github.com/tan7271/OpenPiRoboEval.git#subdirectory=packages/openpi-client --no-cache-dir --no-deps
+echo "Installing openpi-client in openpi_env..."
+conda run -n openpi_env pip install git+https://${GH_TOKEN}@github.com/tan7271/OpenPiRoboEval.git#subdirectory=packages/openpi-client --no-cache-dir --no-deps
 echo "openpi-client installed successfully"
 
-# Install OpenPI
-pip install git+https://${GH_TOKEN}@github.com/tan7271/OpenPiRoboEval.git --no-cache-dir --no-deps --force-reinstall
+echo "Installing OpenPI in openpi_env..."
+conda run -n openpi_env pip install git+https://${GH_TOKEN}@github.com/tan7271/OpenPiRoboEval.git --no-cache-dir --no-deps --force-reinstall
 echo "OpenPI installed successfully"
 
-# Install OpenVLA from source
-echo "Installing OpenVLA from openvla/openvla..."
-pip install git+https://github.com/openvla/openvla.git --no-cache-dir
-echo "OpenVLA installed successfully"
+# Copy RoboEval to openpi_env
+OPENPI_SITE=$(conda run -n openpi_env python -c "import site; print(site.getsitepackages()[0])")
+cp -r $SITE_PACKAGES/roboeval* $OPENPI_SITE/ || true
+cp -r $SITE_PACKAGES/thirdparty $OPENPI_SITE/ || true
+echo "OpenPI environment ready"
 
-echo "===== All dependencies installed ====="
+# Create OpenVLA environment (TEMPORARILY DISABLED - uncomment to enable)
+# echo ""
+# echo "===== Creating OpenVLA Environment ====="
+# conda create -n openvla_env python=3.10 -y
+# 
+# # OpenVLA requires older PyTorch versions
+# echo "Installing OpenVLA-compatible PyTorch in openvla_env..."
+# conda run -n openvla_env pip install torch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 --index-url https://download.pytorch.org/whl/cu121 --no-cache-dir
+# echo "PyTorch installed successfully"
+# 
+# echo "Installing OpenVLA dependencies in openvla_env..."
+# conda run -n openvla_env pip install -r requirements_openvla.txt --no-cache-dir
+# echo "Dependencies installed successfully"
+# 
+# # Install OpenVLA from GitHub
+# echo "Installing OpenVLA from openvla/openvla..."
+# conda run -n openvla_env pip install git+https://github.com/openvla/openvla.git --no-cache-dir
+# echo "OpenVLA installed successfully"
+# 
+# # Copy RoboEval to openvla_env
+# OPENVLA_SITE=$(conda run -n openvla_env python -c "import site; print(site.getsitepackages()[0])")
+# cp -r $SITE_PACKAGES/roboeval* $OPENVLA_SITE/ || true
+# cp -r $SITE_PACKAGES/thirdparty $OPENVLA_SITE/ || true
+# echo "OpenVLA environment ready"
 
+echo ""
+echo "===== Setup Complete ====="
+echo "✓ Base environment: Gradio + RoboEval"
+echo "✓ openpi_env: PyTorch 2.7 + OpenPI"
+echo "ℹ openvla_env: Disabled (uncomment in setup.sh to enable)"