import asyncio import contextlib import logging import time from collections import deque import numpy as np from transport_server_client import RoboticsConsumer, RoboticsProducer from transport_server_client.video import VideoConsumer, VideoProducer from inference_server.models import get_inference_engine, list_supported_policies from inference_server.models.joint_config import JointConfig logger = logging.getLogger(__name__) def busy_wait(seconds): """ Precise timing function for consistent control loops. On some systems, asyncio.sleep is not accurate enough for control loops, so we use busy waiting for short delays. """ if seconds > 0: end_time = asyncio.get_event_loop().time() + seconds while asyncio.get_event_loop().time() < end_time: pass class InferenceSession: """ A single inference session managing one model and its Arena connections. Handles joint values in NORMALIZED VALUES throughout the pipeline. Supports multiple camera streams with different camera names. Supports multiple policy types: ACT, Pi0, SmolVLA, Diffusion Policy. """ def __init__( self, session_id: str, policy_path: str, camera_names: list[str], arena_server_url: str, workspace_id: str, camera_room_ids: dict[str, str], joint_input_room_id: str, joint_output_room_id: str, policy_type: str = "act", language_instruction: str | None = None, ): self.session_id = session_id self.policy_path = policy_path self.policy_type = policy_type.lower() self.camera_names = camera_names self.arena_server_url = arena_server_url self.language_instruction = language_instruction # Validate policy type if self.policy_type not in list_supported_policies(): supported = list_supported_policies() msg = f"Unsupported policy type: {policy_type}. Supported: {supported}" raise ValueError(msg) # Workspace and Room IDs self.workspace_id = workspace_id self.camera_room_ids = camera_room_ids # {camera_name: room_id} self.joint_input_room_id = joint_input_room_id self.joint_output_room_id = joint_output_room_id # Arena clients - multiple camera consumers self.camera_consumers: dict[str, VideoConsumer] = {} # {camera_name: consumer} self.joint_input_consumer: RoboticsConsumer | None = None self.joint_output_producer: RoboticsProducer | None = None # Generic inference engine (supports all policy types) self.inference_engine = None # Session state self.status = "initializing" self.error_message: str | None = None self.inference_task: asyncio.Task | None = None # Data buffers - all in normalized values self.latest_images: dict[str, np.ndarray] = {} # {camera_name: image} self.latest_joint_positions: np.ndarray | None = None # Complete joint state (always 6 joints) - initialized with zeros self.complete_joint_state: np.ndarray = np.zeros(6, dtype=np.float32) self.images_updated: dict[str, bool] = dict.fromkeys(camera_names, False) self.joints_updated = False # Action queue for proper chunking (important for ACT, optional for others) self.action_queue: deque = deque(maxlen=100) # Adjust maxlen as needed self.n_action_steps = 10 # How many actions to use from each chunk # Memory optimization: Clear old actions periodically self.last_queue_cleanup = time.time() self.queue_cleanup_interval = 10.0 # seconds # Control frequency configuration (matches LeRobot defaults) self.control_frequency_hz = 20 # Hz - reduced from 30 to improve performance self.inference_frequency_hz = 2 # Hz - reduced from 3 to improve performance # Statistics self.stats = { "inference_count": 0, "images_received": dict.fromkeys(camera_names, 0), "joints_received": 0, "commands_sent": 0, "errors": 0, "actions_in_queue": 0, "policy_type": self.policy_type, } # Robot responsiveness tracking self.last_command_values: np.ndarray | None = None self.last_joint_check_time = time.time() # Session timeout management self.last_activity_time = time.time() self.timeout_seconds = 600 # 10 minutes self.timeout_check_task: asyncio.Task | None = None async def initialize(self): """Initialize the session by loading the model and setting up Arena connections.""" logger.info( f"Initializing session {self.session_id} with policy type: {self.policy_type}, " f"cameras: {self.camera_names}" ) # Initialize inference engine based on policy type engine_kwargs = { "policy_path": self.policy_path, "camera_names": self.camera_names, } # Add language instruction for policies that support it if ( self.policy_type in {"pi0", "pi0fast", "smolvla"} and self.language_instruction ): engine_kwargs["language_instruction"] = self.language_instruction self.inference_engine = get_inference_engine(self.policy_type, **engine_kwargs) # Load the policy await self.inference_engine.load_policy() # Create Arena clients for each camera for camera_name in self.camera_names: self.camera_consumers[camera_name] = VideoConsumer(self.arena_server_url) self.joint_input_consumer = RoboticsConsumer(self.arena_server_url) self.joint_output_producer = RoboticsProducer(self.arena_server_url) # Set up callbacks self._setup_callbacks() # Connect to rooms await self._connect_to_rooms() # Start receiving video frames from all cameras for camera_name, consumer in self.camera_consumers.items(): await consumer.start_receiving() logger.info(f"Started receiving frames for camera: {camera_name}") # Start timeout monitoring self.timeout_check_task = asyncio.create_task(self._timeout_monitor()) self.status = "ready" logger.info( f"✅ Session {self.session_id} initialized successfully with {self.policy_type} policy " f"and {len(self.camera_names)} cameras" ) def _setup_callbacks(self): """Set up callbacks for Arena clients.""" def create_frame_callback(camera_name: str): """Create a frame callback for a specific camera.""" def on_frame_received(frame_data): """Handle incoming camera frame from VideoConsumer.""" try: metadata = frame_data.metadata width = metadata.get("width", 0) height = metadata.get("height", 0) format_type = metadata.get("format", "rgb24") if format_type == "rgb24" and width > 0 and height > 0: # Convert bytes to numpy array (server sends RGB format) frame_bytes = frame_data.data # Validate frame data size expected_size = height * width * 3 if len(frame_bytes) != expected_size: logger.warning( f"Frame size mismatch for camera {camera_name}: " f"expected {expected_size}, got {len(frame_bytes)}. Skipping frame." ) self.stats["errors"] += 1 return img_rgb = np.frombuffer(frame_bytes, dtype=np.uint8).reshape(( height, width, 3, )) # Store as latest image for inference self.latest_images[camera_name] = img_rgb self.images_updated[camera_name] = True self.stats["images_received"][camera_name] += 1 # Update activity time self.last_activity_time = time.time() else: logger.debug( f"Skipping invalid frame for camera {camera_name}: {format_type}, " f"{width}x{height}" ) except Exception as e: logger.exception( f"Error processing frame for camera {camera_name}: {e}" ) self.stats["errors"] += 1 return on_frame_received # Set up frame callbacks for each camera for camera_name in self.camera_names: callback = create_frame_callback(camera_name) self.camera_consumers[camera_name].on_frame_update(callback) def on_joints_received(joints_data): """Handle incoming joint data from RoboticsConsumer.""" try: joint_values = self._parse_joint_data(joints_data) if joint_values: # Update complete joint state with received values for i, value in enumerate(joint_values[:6]): # Ensure max 6 joints self.complete_joint_state[i] = value self.latest_joint_positions = self.complete_joint_state.copy() self.joints_updated = True self.stats["joints_received"] += 1 # Update activity time self.last_activity_time = time.time() except Exception as e: logger.exception(f"Error processing joint data: {e}") self.stats["errors"] += 1 def on_error(error_msg): """Handle Arena client errors.""" logger.error( f"Arena client error in session {self.session_id}: {error_msg}" ) self.error_message = str(error_msg) self.stats["errors"] += 1 # Set callbacks for RoboticsConsumer self.joint_input_consumer.on_joint_update(on_joints_received) self.joint_input_consumer.on_state_sync(on_joints_received) self.joint_input_consumer.on_error(on_error) # Set error callbacks for VideoConsumers for consumer in self.camera_consumers.values(): consumer.on_error(on_error) def _parse_joint_data(self, joints_data) -> list[float]: """ Parse joint data from Arena message. Expected format: dict with joint names as keys and normalized values. All values are already normalized from the training data pipeline. Args: joints_data: Joint data from Arena message Returns: List of 6 normalized joint values in LeRobot standard order """ return JointConfig.parse_joint_data(joints_data, self.policy_type) def _get_joint_index(self, joint_name: str) -> int | None: """ Get the index of a joint in the standard joint order. Args: joint_name: Name of the joint Returns: Index of the joint, or None if not found """ return JointConfig.get_joint_index(joint_name) async def _connect_to_rooms(self): """Connect to all Arena rooms.""" # Connect to camera rooms as consumer for camera_name, consumer in self.camera_consumers.items(): room_id = self.camera_room_ids[camera_name] success = await consumer.connect( self.workspace_id, room_id, f"{self.session_id}-{camera_name}-consumer" ) if not success: msg = f"Failed to connect to camera room for {camera_name}" raise Exception(msg) logger.info( f"Connected to camera room for {camera_name}: {room_id} in workspace {self.workspace_id}" ) # Connect to joint input room as consumer success = await self.joint_input_consumer.connect( self.workspace_id, self.joint_input_room_id, f"{self.session_id}-joint-input-consumer", ) if not success: msg = "Failed to connect to joint input room" raise Exception(msg) # Connect to joint output room as producer success = await self.joint_output_producer.connect( self.workspace_id, self.joint_output_room_id, f"{self.session_id}-joint-output-producer", ) if not success: msg = "Failed to connect to joint output room" raise Exception(msg) logger.info( f"Connected to all rooms for session {self.session_id} in workspace {self.workspace_id}" ) async def start_inference(self): """Start the inference loop.""" if self.status != "ready": msg = f"Session not ready. Current status: {self.status}" raise Exception(msg) self.status = "running" self.inference_task = asyncio.create_task(self._inference_loop()) logger.info(f"Started inference for session {self.session_id}") async def stop_inference(self): """Stop the inference loop.""" if self.inference_task: self.inference_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self.inference_task self.inference_task = None self.status = "stopped" logger.info(f"Stopped inference for session {self.session_id}") async def restart_inference(self): """Restart the inference loop (stop if running, then start).""" logger.info(f"Restarting inference for session {self.session_id}") # Stop current inference if running await self.stop_inference() # Reset internal state for fresh start self._reset_session_state() # Start inference again await self.start_inference() logger.info(f"Successfully restarted inference for session {self.session_id}") def _reset_session_state(self): """Reset session state for restart.""" # Clear action queue self.action_queue.clear() # Reset complete joint state to zeros self.complete_joint_state.fill(0.0) # Reset image update flags for camera_name in self.camera_names: self.images_updated[camera_name] = False self.joints_updated = False # Reset timing self.last_queue_cleanup = time.time() # Reset inference engine state if available if self.inference_engine: self.inference_engine.reset() # Reset some statistics (but keep cumulative counts) self.stats["actions_in_queue"] = 0 logger.info(f"Reset session state for {self.session_id}") async def _timeout_monitor(self): """Monitor session for inactivity timeout.""" while True: try: await asyncio.sleep(60) # Check every minute current_time = time.time() inactive_time = current_time - self.last_activity_time if inactive_time > self.timeout_seconds: logger.warning( f"Session {self.session_id} has been inactive for " f"{inactive_time:.1f} seconds (timeout: {self.timeout_seconds}s). " f"Marking for cleanup." ) self.status = "timeout" # Don't cleanup here - let the session manager handle it break if inactive_time > self.timeout_seconds * 0.8: # Warn at 80% of timeout logger.info( f"Session {self.session_id} inactive for {inactive_time:.1f}s, " f"will timeout in {self.timeout_seconds - inactive_time:.1f}s" ) except asyncio.CancelledError: logger.info(f"Timeout monitor cancelled for session {self.session_id}") break except Exception as e: logger.exception( f"Error in timeout monitor for session {self.session_id}: {e}" ) break def _all_cameras_have_data(self) -> bool: """Check if we have received data from all cameras.""" return all( camera_name in self.latest_images for camera_name in self.camera_names ) async def _inference_loop(self): """Main inference loop that processes incoming data and sends commands.""" logger.info(f"Starting inference loop for session {self.session_id}") logger.info( f"Control frequency: {self.control_frequency_hz} Hz, Inference frequency: {self.inference_frequency_hz} Hz" ) logger.info( f"Waiting for data from {len(self.camera_names)} cameras: {self.camera_names}" ) inference_counter = 0 target_dt = 1.0 / self.control_frequency_hz # Control loop period inference_interval = ( self.control_frequency_hz // self.inference_frequency_hz ) # How many control steps per inference while True: loop_start_time = asyncio.get_event_loop().time() # Check if we have images from all cameras and joint data if ( self._all_cameras_have_data() and self.latest_joint_positions is not None ): # Removed verbose data logging # Only run inference at the specified frequency and when queue is empty should_run_inference = len(self.action_queue) == 0 or ( inference_counter % inference_interval == 0 and len(self.action_queue) < 3 ) if should_run_inference: # Only log inference runs occasionally to reduce overhead if self.stats["inference_count"] % 10 == 0: logger.info( f"Running inference #{self.stats['inference_count']} for session {self.session_id} " f"(queue length: {len(self.action_queue)})" ) try: # Verify joint positions have correct shape before inference if self.latest_joint_positions.shape != (6,): logger.error( f"Invalid joint positions shape: {self.latest_joint_positions.shape}, " f"expected (6,). Values: {self.latest_joint_positions}" ) # Fix the shape by resetting to complete joint state self.latest_joint_positions = ( self.complete_joint_state.copy() ) logger.debug( f"Running inference with joint positions shape: {self.latest_joint_positions.shape}, " f"values: {self.latest_joint_positions}" ) # Prepare inference arguments inference_kwargs = { "images": self.latest_images, "joint_positions": self.latest_joint_positions, } # Add language instruction for vision-language policies if ( self.policy_type in {"pi0", "pi0fast", "smolvla"} and self.language_instruction ): inference_kwargs["task"] = self.language_instruction # Run inference to get action chunk predicted_actions = await self.inference_engine.predict( **inference_kwargs ) # ACT returns a chunk of actions, we need to queue them if len(predicted_actions.shape) == 1: # Single action returned, use it directly actions_to_queue = [predicted_actions] else: # Multiple actions in chunk, take first n_action_steps actions_to_queue = predicted_actions[: self.n_action_steps] # Add actions to queue for action in actions_to_queue: joint_commands = ( self.inference_engine.get_joint_commands_with_names( action ) ) self.action_queue.append(joint_commands) self.stats["inference_count"] += 1 # Reset image update flags for camera_name in self.camera_names: self.images_updated[camera_name] = False self.joints_updated = False logger.debug( f"Added {len(actions_to_queue)} actions to queue for session {self.session_id}" ) except Exception as e: logger.exception( f"Inference failed for session {self.session_id}: {e}" ) self.stats["errors"] += 1 # Send action from queue if available if len(self.action_queue) > 0: joint_commands = self.action_queue.popleft() try: # Only log commands occasionally if self.stats["commands_sent"] % 100 == 0: logger.info( f"🤖 Sent {self.stats['commands_sent']} commands. Latest: {joint_commands[0]['name']}={joint_commands[0]['value']:.1f}" ) await self.joint_output_producer.send_joint_update( joint_commands ) self.stats["commands_sent"] += 1 self.stats["actions_in_queue"] = len(self.action_queue) # Store command values for responsiveness check command_values = np.array( [cmd["value"] for cmd in joint_commands], dtype=np.float32 ) self.last_command_values = command_values except Exception as e: logger.exception( f"Failed to send joint commands for session {self.session_id}: {e}" ) self.stats["errors"] += 1 # Log when queue is empty occasionally elif inference_counter % 100 == 0: logger.debug( f"No actions in queue to send (inference #{inference_counter})" ) # Periodic memory cleanup current_time = asyncio.get_event_loop().time() if current_time - self.last_queue_cleanup > self.queue_cleanup_interval: # Clear stale actions if queue is getting full if len(self.action_queue) > 80: # 80% of maxlen logger.debug( f"Clearing stale actions from queue for session {self.session_id}" ) self.action_queue.clear() self.last_queue_cleanup = current_time inference_counter += 1 # Precise timing control for consistent control frequency elapsed_time = asyncio.get_event_loop().time() - loop_start_time sleep_time = target_dt - elapsed_time if sleep_time > 0.001: # Use asyncio.sleep for longer waits await asyncio.sleep(sleep_time) elif sleep_time > 0: # Use busy_wait for precise short delays busy_wait(sleep_time) elif sleep_time < -0.01: # Log if we're running significantly slow logger.warning( f"Control loop running slow for session {self.session_id}: {elapsed_time * 1000:.1f}ms (target: {target_dt * 1000:.1f}ms)" ) async def cleanup(self): """Clean up session resources.""" logger.info(f"Cleaning up session {self.session_id}") # Stop timeout monitoring if self.timeout_check_task: self.timeout_check_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self.timeout_check_task self.timeout_check_task = None # Stop inference await self.stop_inference() # Disconnect Arena clients for camera_name, consumer in self.camera_consumers.items(): await consumer.stop_receiving() await consumer.disconnect() logger.info(f"Disconnected camera consumer for {camera_name}") if self.joint_input_consumer: await self.joint_input_consumer.disconnect() if self.joint_output_producer: await self.joint_output_producer.disconnect() # Clean up inference engine if self.inference_engine: del self.inference_engine self.inference_engine = None logger.info(f"Session {self.session_id} cleanup completed") def get_status(self) -> dict: """Get current session status.""" status_dict = { "session_id": self.session_id, "status": self.status, "policy_path": self.policy_path, "policy_type": self.policy_type, "camera_names": self.camera_names, "workspace_id": self.workspace_id, "rooms": { "workspace_id": self.workspace_id, "camera_room_ids": self.camera_room_ids, "joint_input_room_id": self.joint_input_room_id, "joint_output_room_id": self.joint_output_room_id, }, "stats": self.stats.copy(), "error_message": self.error_message, "joint_state": { "complete_joint_state": self.complete_joint_state.tolist(), "latest_joint_positions": ( self.latest_joint_positions.tolist() if self.latest_joint_positions is not None else None ), "joint_state_shape": ( self.latest_joint_positions.shape if self.latest_joint_positions is not None else None ), }, } # Add inference engine stats if available if self.inference_engine: status_dict["inference_stats"] = self.inference_engine.get_model_info() return status_dict class SessionManager: """Manages multiple inference sessions and their lifecycle.""" def __init__(self): self.sessions: dict[str, InferenceSession] = {} self.cleanup_task: asyncio.Task | None = None self._start_cleanup_task() def _start_cleanup_task(self): """Start the automatic cleanup task for timed-out sessions.""" try: # Only start if we're in an async context loop = asyncio.get_running_loop() self.cleanup_task = loop.create_task(self._periodic_cleanup()) except RuntimeError: # No event loop running yet, will start later pass async def _periodic_cleanup(self): """Periodically check for and clean up timed-out sessions.""" while True: try: await asyncio.sleep(300) # Check every 5 minutes # Find sessions that have timed out timed_out_sessions = [] for session_id, session in self.sessions.items(): if session.status == "timeout": timed_out_sessions.append(session_id) # Clean up timed-out sessions for session_id in timed_out_sessions: logger.info(f"Auto-cleaning up timed-out session: {session_id}") await self.delete_session(session_id) except asyncio.CancelledError: logger.info("Periodic cleanup task cancelled") break except Exception as e: logger.exception(f"Error in periodic cleanup: {e}") # Continue running even if there's an error async def create_session( self, session_id: str, policy_path: str, camera_names: list[str] | None = None, arena_server_url: str = "http://localhost:8000", workspace_id: str | None = None, policy_type: str = "act", language_instruction: str | None = None, ) -> dict[str, str]: """Create a new inference session.""" if camera_names is None: camera_names = ["front"] if session_id in self.sessions: msg = f"Session {session_id} already exists" raise ValueError(msg) # Create camera rooms using VideoProducer video_temp_client = VideoProducer(arena_server_url) camera_room_ids = {} # Use provided workspace_id or create new one if workspace_id: target_workspace_id = workspace_id logger.info( f"Using provided workspace ID {target_workspace_id} for session {session_id}" ) # Create all camera rooms in the specified workspace for camera_name in camera_names: _, room_id = await video_temp_client.create_room( workspace_id=target_workspace_id, room_id=f"{session_id}-{camera_name}", ) camera_room_ids[camera_name] = room_id else: # Create first camera room to get new workspace_id first_camera = camera_names[0] target_workspace_id, first_room_id = await video_temp_client.create_room( room_id=f"{session_id}-{first_camera}" ) logger.info( f"Generated new workspace ID {target_workspace_id} for session {session_id}" ) # Store the first room camera_room_ids[first_camera] = first_room_id # Create remaining camera rooms in the same workspace for camera_name in camera_names[1:]: _, room_id = await video_temp_client.create_room( workspace_id=target_workspace_id, room_id=f"{session_id}-{camera_name}", ) camera_room_ids[camera_name] = room_id # Create joint rooms using RoboticsProducer in the same workspace robotics_temp_client = RoboticsProducer(arena_server_url) _, joint_input_room_id = await robotics_temp_client.create_room( workspace_id=target_workspace_id, room_id=f"{session_id}-joint-input" ) _, joint_output_room_id = await robotics_temp_client.create_room( workspace_id=target_workspace_id, room_id=f"{session_id}-joint-output" ) logger.info( f"Created rooms for session {session_id} in workspace {target_workspace_id}:" ) for camera_name, room_id in camera_room_ids.items(): logger.info(f" Camera room ({camera_name}): {room_id}") logger.info(f" Joint input room: {joint_input_room_id}") logger.info(f" Joint output room: {joint_output_room_id}") # Create session session = InferenceSession( session_id=session_id, policy_path=policy_path, camera_names=camera_names, arena_server_url=arena_server_url, workspace_id=target_workspace_id, camera_room_ids=camera_room_ids, joint_input_room_id=joint_input_room_id, joint_output_room_id=joint_output_room_id, policy_type=policy_type, language_instruction=language_instruction, ) # Initialize session await session.initialize() # Store session self.sessions[session_id] = session # Start cleanup task if not already running if not self.cleanup_task or self.cleanup_task.done(): self._start_cleanup_task() return { "workspace_id": target_workspace_id, "camera_room_ids": camera_room_ids, "joint_input_room_id": joint_input_room_id, "joint_output_room_id": joint_output_room_id, } async def get_session_status(self, session_id: str) -> dict: """Get status of a specific session.""" if session_id not in self.sessions: msg = f"Session {session_id} not found" raise KeyError(msg) return self.sessions[session_id].get_status() async def start_inference(self, session_id: str): """Start inference for a specific session.""" if session_id not in self.sessions: msg = f"Session {session_id} not found" raise KeyError(msg) await self.sessions[session_id].start_inference() async def stop_inference(self, session_id: str): """Stop inference for a specific session.""" if session_id not in self.sessions: msg = f"Session {session_id} not found" raise KeyError(msg) await self.sessions[session_id].stop_inference() async def restart_inference(self, session_id: str): """Restart inference for a specific session.""" if session_id not in self.sessions: msg = f"Session {session_id} not found" raise KeyError(msg) await self.sessions[session_id].restart_inference() async def delete_session(self, session_id: str): """Delete a session and clean up all resources.""" if session_id not in self.sessions: msg = f"Session {session_id} not found" raise KeyError(msg) session = self.sessions[session_id] await session.cleanup() del self.sessions[session_id] logger.info(f"Deleted session {session_id}") async def list_sessions(self) -> list[dict]: """List all sessions with their status.""" return [session.get_status() for session in self.sessions.values()] async def cleanup_all_sessions(self): """Clean up all sessions.""" logger.info("Cleaning up all sessions...") # Stop the cleanup task if self.cleanup_task: self.cleanup_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self.cleanup_task self.cleanup_task = None # Clean up all sessions for session_id in list(self.sessions.keys()): await self.delete_session(session_id) logger.info("All sessions cleaned up")