Spaces:
Sleeping
Sleeping
| import logging | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from lerobot.common.policies.act.modeling_act import ACTPolicy | |
| from lerobot.common.utils.utils import init_logging | |
| from .base_inference import BaseInferenceEngine | |
| logger = logging.getLogger(__name__) | |
| class ACTInferenceEngine(BaseInferenceEngine): | |
| """ | |
| ACT (Action Chunking Transformer) inference engine. | |
| Handles image preprocessing, joint normalization, and action prediction | |
| for ACT models with proper action chunking. | |
| """ | |
| def __init__( | |
| self, | |
| policy_path: str, | |
| camera_names: list[str], | |
| use_custom_joint_names: bool = True, | |
| device: str | None = None, | |
| ): | |
| super().__init__(policy_path, camera_names, use_custom_joint_names, device) | |
| # ACT-specific configuration | |
| self.chunk_size = 10 # Default chunk size for ACT | |
| self.action_history = [] # Store recent actions for chunking | |
| async def load_policy(self): | |
| """Load the ACT policy from the specified path.""" | |
| logger.info(f"Loading ACT policy from: {self.policy_path}") | |
| try: | |
| # Initialize hydra config for LeRobot | |
| init_logging() | |
| # Load the ACT policy | |
| self.policy = ACTPolicy.from_pretrained(self.policy_path) | |
| self.policy.to(self.device) | |
| self.policy.eval() | |
| # Set up image transforms based on policy config | |
| if hasattr(self.policy, "config"): | |
| self._setup_image_transforms() | |
| self.is_loaded = True | |
| logger.info(f"✅ ACT policy loaded successfully on {self.device}") | |
| except Exception as e: | |
| logger.exception(f"Failed to load ACT policy from {self.policy_path}") | |
| msg = f"Failed to load ACT policy: {e}" | |
| raise RuntimeError(msg) from e | |
| def _setup_image_transforms(self): | |
| """Set up image transforms based on the policy configuration.""" | |
| try: | |
| # Get image size from policy config | |
| config = self.policy.config | |
| image_size = getattr(config, "image_size", 224) | |
| # Create transforms for each camera | |
| for camera_name in self.camera_names: | |
| # Use policy-specific transforms if available | |
| if hasattr(self.policy, "image_processor"): | |
| # Use the policy's image processor | |
| self.image_transforms[camera_name] = self.policy.image_processor | |
| else: | |
| # Fall back to default transform with correct size | |
| from torchvision import transforms | |
| self.image_transforms[camera_name] = transforms.Compose([ | |
| transforms.Resize((image_size, image_size)), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] | |
| ), | |
| ]) | |
| except Exception as e: | |
| logger.warning(f"Could not set up image transforms: {e}. Using defaults.") | |
| async def predict( | |
| self, images: dict[str, np.ndarray], joint_positions: np.ndarray, **kwargs | |
| ) -> np.ndarray: | |
| """ | |
| Run ACT inference to predict actions. | |
| Args: | |
| images: Dictionary of {camera_name: rgb_image_array} | |
| joint_positions: Current joint positions in LeRobot standard order | |
| **kwargs: Additional arguments (unused for ACT) | |
| Returns: | |
| Array of predicted actions (chunk of actions for ACT) | |
| """ | |
| if not self.is_loaded: | |
| msg = "Policy not loaded. Call load_policy() first." | |
| raise RuntimeError(msg) | |
| try: | |
| # Preprocess inputs | |
| processed_images = self.preprocess_images(images) | |
| processed_joints = self.preprocess_joint_positions(joint_positions) | |
| # Prepare batch inputs for ACT | |
| batch = self._prepare_batch(processed_images, processed_joints) | |
| # Run inference | |
| with torch.no_grad(): | |
| # ACT returns a chunk of actions | |
| action_chunk = self.policy.predict(batch) | |
| # Convert to numpy | |
| if isinstance(action_chunk, torch.Tensor): | |
| action_chunk = action_chunk.cpu().numpy() | |
| # Store in action history | |
| self.action_history.append(action_chunk) | |
| if len(self.action_history) > 10: # Keep last 10 chunks | |
| self.action_history.pop(0) | |
| logger.debug(f"ACT predicted action chunk shape: {action_chunk.shape}") | |
| return action_chunk | |
| except Exception as e: | |
| logger.exception("ACT inference failed") | |
| msg = f"ACT inference failed: {e}" | |
| raise RuntimeError(msg) from e | |
| def _prepare_batch( | |
| self, images: dict[str, torch.Tensor], joints: torch.Tensor | |
| ) -> dict: | |
| """ | |
| Prepare batch inputs for ACT model. | |
| Args: | |
| images: Preprocessed images | |
| joints: Preprocessed joint positions | |
| Returns: | |
| Batch dictionary for ACT model | |
| """ | |
| batch = {} | |
| # Add images to batch | |
| for camera_name, image_tensor in images.items(): | |
| # Add batch dimension if needed | |
| if len(image_tensor.shape) == 3: | |
| image_tensor = image_tensor.unsqueeze(0) | |
| batch[f"observation.images.{camera_name}"] = image_tensor | |
| # Add joint positions | |
| if len(joints.shape) == 1: | |
| joints = joints.unsqueeze(0) | |
| batch["observation.state"] = joints | |
| return batch | |
| def reset(self): | |
| """Reset ACT-specific state.""" | |
| super().reset() | |
| self.action_history = [] | |
| # Reset ACT model state if it has one | |
| if self.policy and hasattr(self.policy, "reset"): | |
| self.policy.reset() | |
| def get_model_info(self) -> dict: | |
| """Get ACT-specific model information.""" | |
| info = super().get_model_info() | |
| info.update({ | |
| "policy_type": "act", | |
| "chunk_size": self.chunk_size, | |
| "action_history_length": len(self.action_history), | |
| }) | |
| return info | |
| # Utility functions for data transformation | |
| def image_bgr_to_rgb(image: np.ndarray) -> np.ndarray: | |
| """Convert BGR image to RGB (useful for OpenCV cameras).""" | |
| return cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| def resize_image(image: np.ndarray, target_size: tuple[int, int]) -> np.ndarray: | |
| """Resize image to target size (width, height).""" | |
| return cv2.resize(image, target_size) | |