profile

examples/rtc/eval_with_real_robot_profiled.py

@@ -0,0 +1,631 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Profiled version of eval_with_real_robot.py for performance analysis.
+
+This version adds detailed timing measurements for:
+- Policy inference
+- Preprocessing
+- Postprocessing
+- Action queue operations
+- Robot communication
+- Thread execution times
+
+Usage: Same as eval_with_real_robot.py but with profiling output.
+"""
+
+import logging
+import math
+import sys
+import time
+import traceback
+from collections import defaultdict
+from dataclasses import dataclass, field
+from threading import Event, Lock, Thread
+
+import torch
+from torch import Tensor
+
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
+from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
+from lerobot.configs import parser
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.types import RTCAttentionSchedule
+from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
+from lerobot.policies.factory import get_policy_class, make_pre_post_processors
+from lerobot.policies.rtc.action_queue import ActionQueue
+from lerobot.policies.rtc.configuration_rtc import RTCConfig
+from lerobot.policies.rtc.latency_tracker import LatencyTracker
+from lerobot.processor.factory import (
+    make_default_robot_action_processor,
+    make_default_robot_observation_processor,
+)
+from lerobot.rl.process import ProcessSignalHandler
+from lerobot.robots import (  # noqa: F401
+    Robot,
+    RobotConfig,
+    koch_follower,
+    so100_follower,
+    so101_follower,
+)
+from lerobot.robots.utils import make_robot_from_config
+from lerobot.utils.constants import OBS_IMAGES
+from lerobot.utils.hub import HubMixin
+from lerobot.utils.utils import init_logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class ProfileTimer:
+    """Context manager and utility class for timing code sections."""
+
+    def __init__(self, name: str, stats_dict: dict):
+        self.name = name
+        self.stats_dict = stats_dict
+        self.start_time = None
+
+    def __enter__(self):
+        self.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, *args):
+        elapsed = time.perf_counter() - self.start_time
+        if self.name not in self.stats_dict:
+            self.stats_dict[self.name] = []
+        self.stats_dict[self.name].append(elapsed)
+
+
+class ProfilingStats:
+    """Global profiling statistics collector."""
+
+    def __init__(self):
+        self.stats = defaultdict(list)
+        self.lock = Lock()
+
+    def record(self, name: str, duration: float):
+        with self.lock:
+            self.stats[name].append(duration)
+
+    def timer(self, name: str):
+        """Return a context manager for timing."""
+        return ProfileTimer(name, self.stats)
+
+    def get_summary(self) -> dict[str, dict[str, float]]:
+        """Get summary statistics for all timings."""
+        with self.lock:
+            summary = {}
+            for name, times in self.stats.items():
+                if times:
+                    summary[name] = {
+                        "count": len(times),
+                        "mean": sum(times) / len(times),
+                        "min": min(times),
+                        "max": max(times),
+                        "total": sum(times),
+                    }
+            return summary
+
+    def print_summary(self):
+        """Print formatted summary of all timings."""
+        summary = self.get_summary()
+        
+        logger.info("\n" + "=" * 80)
+        logger.info("PROFILING SUMMARY")
+        logger.info("=" * 80)
+        
+        # Sort by total time (descending)
+        sorted_items = sorted(summary.items(), key=lambda x: x[1]["total"], reverse=True)
+        
+        for name, stats in sorted_items:
+            logger.info(f"\n{name}:")
+            logger.info(f"  Count:     {stats['count']}")
+            logger.info(f"  Mean:      {stats['mean']*1000:.2f} ms")
+            logger.info(f"  Min:       {stats['min']*1000:.2f} ms")
+            logger.info(f"  Max:       {stats['max']*1000:.2f} ms")
+            logger.info(f"  Total:     {stats['total']:.2f} s")
+            logger.info(f"  Hz:        {stats['count']/stats['total']:.2f}")
+        
+        logger.info("\n" + "=" * 80)
+
+
+# Global profiling stats
+profiling_stats = ProfilingStats()
+
+
+class RobotWrapper:
+    def __init__(self, robot: Robot):
+        self.robot = robot
+        self.lock = Lock()
+
+    def get_observation(self) -> dict[str, Tensor]:
+        with profiling_stats.timer("robot.get_observation"):
+            with self.lock:
+                return self.robot.get_observation()
+
+    def send_action(self, action: Tensor):
+        with profiling_stats.timer("robot.send_action"):
+            with self.lock:
+                self.robot.send_action(action)
+
+    def observation_features(self) -> list[str]:
+        with self.lock:
+            return self.robot.observation_features
+
+    def action_features(self) -> list[str]:
+        with self.lock:
+            return self.robot.action_features
+
+
+@dataclass
+class RTCDemoConfig(HubMixin):
+    """Configuration for RTC demo with action chunking policies and real robots."""
+
+    # Policy configuration
+    policy: PreTrainedConfig | None = None
+
+    # Robot configuration
+    robot: RobotConfig | None = None
+
+    # RTC configuration
+    rtc: RTCConfig = field(
+        default_factory=lambda: RTCConfig(
+            execution_horizon=10,
+            max_guidance_weight=1.0,
+            prefix_attention_schedule=RTCAttentionSchedule.EXP,
+        )
+    )
+
+    # Demo parameters
+    duration: float = 30.0  # Duration to run the demo (seconds)
+    fps: float = 10.0  # Action execution frequency (Hz)
+
+    # Compute device
+    device: str | None = None  # Device to run on (cuda, cpu, auto)
+
+    # Get new actions horizon. The amount of executed steps after which will be requested new actions.
+    # It should be higher than inference delay + execution horizon.
+    action_queue_size_to_get_new_actions: int = 30
+
+    # Task to execute
+    task: str = field(default="", metadata={"help": "Task to execute"})
+
+    # Torch compile configuration
+    use_torch_compile: bool = field(
+        default=False,
+        metadata={"help": "Use torch.compile for faster inference (PyTorch 2.0+)"},
+    )
+
+    torch_compile_backend: str = field(
+        default="inductor",
+        metadata={"help": "Backend for torch.compile (inductor, aot_eager, cudagraphs)"},
+    )
+
+    torch_compile_mode: str = field(
+        default="default",
+        metadata={"help": "Compilation mode (default, reduce-overhead, max-autotune)"},
+    )
+
+    torch_compile_disable_cudagraphs: bool = field(
+        default=True,
+        metadata={
+            "help": "Disable CUDA graphs in torch.compile. Required due to in-place tensor "
+            "operations in denoising loop (x_t += dt * v_t) which cause tensor aliasing issues."
+        },
+    )
+
+    def __post_init__(self):
+        # HACK: We parse again the cli args here to get the pretrained path if there was one.
+        policy_path = parser.get_path_arg("policy")
+        if policy_path:
+            cli_overrides = parser.get_cli_overrides("policy")
+            self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides)
+            self.policy.pretrained_path = policy_path
+        else:
+            raise ValueError("Policy path is required")
+
+        # Validate that robot configuration is provided
+        if self.robot is None:
+            raise ValueError("Robot configuration must be provided")
+
+    @classmethod
+    def __get_path_fields__(cls) -> list[str]:
+        """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
+        return ["policy"]
+
+
+def is_image_key(k: str) -> bool:
+    return k.startswith(OBS_IMAGES)
+
+
+def get_actions(
+    policy,
+    robot: RobotWrapper,
+    robot_observation_processor,
+    action_queue: ActionQueue,
+    shutdown_event: Event,
+    cfg: RTCDemoConfig,
+):
+    """Thread function to request action chunks from the policy with profiling.
+
+    Args:
+        policy: The policy instance (SmolVLA, Pi0, etc.)
+        robot: The robot instance for getting observations
+        robot_observation_processor: Processor for raw robot observations
+        action_queue: Queue to put new action chunks
+        shutdown_event: Event to signal shutdown
+        cfg: Demo configuration
+    """
+    try:
+        logger.info("[GET_ACTIONS] Starting get actions thread")
+
+        latency_tracker = LatencyTracker()  # Track latency of action chunks
+        fps = cfg.fps
+        time_per_chunk = 1.0 / fps
+
+        dataset_features = hw_to_dataset_features(robot.observation_features(), "observation")
+        policy_device = policy.config.device
+
+        # Load preprocessor and postprocessor from pretrained files
+        logger.info(f"[GET_ACTIONS] Loading preprocessor/postprocessor from {cfg.policy.pretrained_path}")
+
+        preprocessor, postprocessor = make_pre_post_processors(
+            policy_cfg=cfg.policy,
+            pretrained_path=cfg.policy.pretrained_path,
+            dataset_stats=None,  # Will load from pretrained processor files
+            preprocessor_overrides={
+                "device_processor": {"device": cfg.policy.device},
+            },
+        )
+
+        logger.info("[GET_ACTIONS] Preprocessor/postprocessor loaded successfully with embedded stats")
+
+        get_actions_threshold = cfg.action_queue_size_to_get_new_actions
+
+        if not cfg.rtc.enabled:
+            get_actions_threshold = 0
+
+        inference_count = 0
+
+        while not shutdown_event.is_set():
+            if action_queue.qsize() <= get_actions_threshold:
+                with profiling_stats.timer("get_actions.total_iteration"):
+                    inference_count += 1
+                    logger.info(f"[GET_ACTIONS] Starting inference #{inference_count}")
+
+                    current_time = time.perf_counter()
+                    action_index_before_inference = action_queue.get_action_index()
+                    
+                    with profiling_stats.timer("get_actions.get_prev_actions"):
+                        prev_actions = action_queue.get_left_over()
+
+                    inference_latency = latency_tracker.max()
+                    inference_delay = math.ceil(inference_latency / time_per_chunk)
+
+                    # Get observation
+                    obs = robot.get_observation()
+
+                    # Apply robot observation processor
+                    with profiling_stats.timer("get_actions.robot_obs_processing"):
+                        obs_processed = robot_observation_processor(obs)
+
+                    # Build dataset frame
+                    with profiling_stats.timer("get_actions.build_dataset_frame"):
+                        obs_with_policy_features = build_dataset_frame(
+                            dataset_features, obs_processed, prefix="observation"
+                        )
+
+                    # Convert to tensors and normalize
+                    with profiling_stats.timer("get_actions.tensor_conversion"):
+                        for name in obs_with_policy_features:
+                            obs_with_policy_features[name] = torch.from_numpy(obs_with_policy_features[name])
+                            if "image" in name:
+                                obs_with_policy_features[name] = (
+                                    obs_with_policy_features[name].type(torch.float32) / 255
+                                )
+                                obs_with_policy_features[name] = (
+                                    obs_with_policy_features[name].permute(2, 0, 1).contiguous()
+                                )
+                            obs_with_policy_features[name] = obs_with_policy_features[name].unsqueeze(0)
+                            obs_with_policy_features[name] = obs_with_policy_features[name].to(policy_device)
+
+                        obs_with_policy_features["task"] = [cfg.task]
+                        obs_with_policy_features["robot_type"] = (
+                            robot.robot.name if hasattr(robot.robot, "name") else ""
+                        )
+
+                    # Preprocessing
+                    with profiling_stats.timer("get_actions.preprocessing"):
+                        preproceseded_obs = preprocessor(obs_with_policy_features)
+
+                    # Policy inference
+                    with profiling_stats.timer("get_actions.policy_inference"):
+                        actions = policy.predict_action_chunk(
+                            preproceseded_obs,
+                            inference_delay=inference_delay,
+                            prev_chunk_left_over=prev_actions,
+                        )
+
+                    # Clone for RTC
+                    with profiling_stats.timer("get_actions.clone_actions"):
+                        original_actions = actions.squeeze(0).clone()
+
+                    # Postprocessing
+                    with profiling_stats.timer("get_actions.postprocessing"):
+                        postprocessed_actions = postprocessor(actions)
+                        postprocessed_actions = postprocessed_actions.squeeze(0)
+
+                    # Update latency tracker
+                    new_latency = time.perf_counter() - current_time
+                    new_delay = math.ceil(new_latency / time_per_chunk)
+                    latency_tracker.add(new_latency)
+
+                    logger.info(
+                        f"[GET_ACTIONS] Inference #{inference_count} completed in {new_latency*1000:.2f}ms "
+                        f"(delay={new_delay} chunks)"
+                    )
+
+                    if cfg.action_queue_size_to_get_new_actions < cfg.rtc.execution_horizon + new_delay:
+                        logger.warning(
+                            "[GET_ACTIONS] cfg.action_queue_size_to_get_new_actions Too small, "
+                            "It should be higher than inference delay + execution horizon."
+                        )
+
+                    # Merge into action queue
+                    with profiling_stats.timer("get_actions.action_queue_merge"):
+                        action_queue.merge(
+                            original_actions, postprocessed_actions, new_delay, action_index_before_inference
+                        )
+            else:
+                # Small sleep to prevent busy waiting
+                time.sleep(0.1)
+
+        logger.info("[GET_ACTIONS] get actions thread shutting down")
+    except Exception as e:
+        logger.error(f"[GET_ACTIONS] Fatal exception in get_actions thread: {e}")
+        logger.error(traceback.format_exc())
+        sys.exit(1)
+
+
+def actor_control(
+    robot: RobotWrapper,
+    robot_action_processor,
+    action_queue: ActionQueue,
+    shutdown_event: Event,
+    cfg: RTCDemoConfig,
+):
+    """Thread function to execute actions on the robot with profiling.
+
+    Args:
+        robot: The robot instance
+        action_queue: Queue to get actions from
+        shutdown_event: Event to signal shutdown
+        cfg: Demo configuration
+    """
+    try:
+        logger.info("[ACTOR] Starting actor thread")
+
+        action_count = 0
+        action_interval = 1.0 / cfg.fps
+
+        while not shutdown_event.is_set():
+            start_time = time.perf_counter()
+
+            with profiling_stats.timer("actor.total_iteration"):
+                # Get action from queue
+                with profiling_stats.timer("actor.queue_get"):
+                    action = action_queue.get()
+
+                if action is not None:
+                    # Process action
+                    with profiling_stats.timer("actor.action_processing"):
+                        action = action.cpu()
+                        action_dict = {key: action[i].item() for i, key in enumerate(robot.action_features())}
+                        action_processed = robot_action_processor((action_dict, None))
+                    
+                    # Send to robot (includes robot.send_action timing)
+                    robot.send_action(action_processed)
+                    action_count += 1
+
+            # Sleep to maintain target FPS
+            dt_s = time.perf_counter() - start_time
+            sleep_time = max(0, (action_interval - dt_s) - 0.001)
+            if sleep_time > 0:
+                time.sleep(sleep_time)
+
+        logger.info(f"[ACTOR] Actor thread shutting down. Total actions executed: {action_count}")
+    except Exception as e:
+        logger.error(f"[ACTOR] Fatal exception in actor_control thread: {e}")
+        logger.error(traceback.format_exc())
+        sys.exit(1)
+
+
+def _apply_torch_compile(policy, cfg: RTCDemoConfig):
+    """Apply torch.compile to the policy's predict_action_chunk method.
+
+    Args:
+        policy: Policy instance to compile
+        cfg: Configuration containing torch compile settings
+
+    Returns:
+        Policy with compiled predict_action_chunk method
+    """
+
+    # PI models handle their own compilation
+    if policy.type == "pi05" or policy.type == "pi0":
+        return policy
+
+    try:
+        # Check if torch.compile is available (PyTorch 2.0+)
+        if not hasattr(torch, "compile"):
+            logger.warning(
+                f"torch.compile is not available. Requires PyTorch 2.0+. "
+                f"Current version: {torch.__version__}. Skipping compilation."
+            )
+            return policy
+
+        logger.info("Applying torch.compile to predict_action_chunk...")
+        logger.info(f"  Backend: {cfg.torch_compile_backend}")
+        logger.info(f"  Mode: {cfg.torch_compile_mode}")
+        logger.info(f"  Disable CUDA graphs: {cfg.torch_compile_disable_cudagraphs}")
+
+        # Compile the predict_action_chunk method
+        compile_kwargs = {
+            "backend": cfg.torch_compile_backend,
+            "mode": cfg.torch_compile_mode,
+        }
+
+        # Disable CUDA graphs if requested (prevents tensor aliasing issues)
+        if cfg.torch_compile_disable_cudagraphs:
+            compile_kwargs["options"] = {"triton.cudagraphs": False}
+
+        original_method = policy.predict_action_chunk
+        compiled_method = torch.compile(original_method, **compile_kwargs)
+        policy.predict_action_chunk = compiled_method
+        logger.info("✓ Successfully compiled predict_action_chunk")
+
+    except Exception as e:
+        logger.error(f"Failed to apply torch.compile: {e}")
+        logger.warning("Continuing without torch.compile")
+
+    return policy
+
+
+@parser.wrap()
+def demo_cli(cfg: RTCDemoConfig):
+    """Main entry point for RTC demo with profiling."""
+
+    # Initialize logging
+    init_logging()
+
+    logger.info(f"Using device: {cfg.device}")
+    logger.info("=" * 80)
+    logger.info("PROFILING MODE ENABLED")
+    logger.info("=" * 80)
+
+    # Setup signal handler for graceful shutdown
+    signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
+    shutdown_event = signal_handler.shutdown_event
+
+    policy = None
+    robot = None
+    get_actions_thread = None
+    actor_thread = None
+
+    policy_class = get_policy_class(cfg.policy.type)
+
+    # Load config and set compile_model for pi0/pi05 models
+    config = PreTrainedConfig.from_pretrained(cfg.policy.pretrained_path)
+
+    if cfg.policy.type == "pi05" or cfg.policy.type == "pi0":
+        config.compile_model = cfg.use_torch_compile
+
+    policy = policy_class.from_pretrained(cfg.policy.pretrained_path, config=config)
+
+    # Turn on RTC
+    policy.config.rtc_config = cfg.rtc
+
+    # Init RTC processor
+    policy.init_rtc_processor()
+
+    assert policy.name in ["smolvla", "pi05", "pi0"], "Only smolvla, pi05, and pi0 are supported for RTC"
+
+    policy = policy.to(cfg.device)
+    policy.eval()
+
+    # Apply torch.compile to predict_action_chunk method if enabled
+    if cfg.use_torch_compile:
+        policy = _apply_torch_compile(policy, cfg)
+
+    # Create robot
+    logger.info(f"Initializing robot: {cfg.robot.type}")
+    robot = make_robot_from_config(cfg.robot)
+    robot.connect()
+    robot_wrapper = RobotWrapper(robot)
+
+    # Create robot observation processor
+    robot_observation_processor = make_default_robot_observation_processor()
+    robot_action_processor = make_default_robot_action_processor()
+
+    # Create action queue for communication between threads
+    action_queue = ActionQueue(cfg.rtc)
+
+    # Start chunk requester thread
+    get_actions_thread = Thread(
+        target=get_actions,
+        args=(policy, robot_wrapper, robot_observation_processor, action_queue, shutdown_event, cfg),
+        daemon=True,
+        name="GetActions",
+    )
+    get_actions_thread.start()
+    logger.info("Started get actions thread")
+
+    # Start action executor thread
+    actor_thread = Thread(
+        target=actor_control,
+        args=(robot_wrapper, robot_action_processor, action_queue, shutdown_event, cfg),
+        daemon=True,
+        name="Actor",
+    )
+    actor_thread.start()
+    logger.info("Started actor thread")
+
+    logger.info("Started stop by duration thread")
+
+    # Main thread monitors for duration or shutdown
+    logger.info(f"Running demo for {cfg.duration} seconds...")
+    start_time = time.time()
+
+    while not shutdown_event.is_set() and (time.time() - start_time) < cfg.duration:
+        time.sleep(10)
+
+        # Log queue status periodically
+        if int(time.time() - start_time) % 5 == 0:
+            logger.info(f"[MAIN] Action queue size: {action_queue.qsize()}")
+
+        if time.time() - start_time > cfg.duration:
+            break
+
+    logger.info("Demo duration reached or shutdown requested")
+
+    # Signal shutdown
+    shutdown_event.set()
+
+    # Wait for threads to finish
+    if get_actions_thread and get_actions_thread.is_alive():
+        logger.info("Waiting for chunk requester thread to finish...")
+        get_actions_thread.join()
+
+    if actor_thread and actor_thread.is_alive():
+        logger.info("Waiting for action executor thread to finish...")
+        actor_thread.join()
+
+    # Cleanup robot
+    if robot:
+        robot.disconnect()
+        logger.info("Robot disconnected")
+
+    # Print profiling summary
+    profiling_stats.print_summary()
+
+    logger.info("Cleanup completed")
+
+
+if __name__ == "__main__":
+    demo_cli()
+    logging.info("RTC demo finished")
+