# Copyright (C) 2025-2026 LejuRobotics. # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . # # --- # # This project includes code from LeRobot (https://github.com/huggingface/lerobot), # which is licensed under the Apache License, Version 2.0. """ This script demonstrates how to evaluate a pretrained policy from the HuggingFace Hub or from your local training outputs directory. In the latter case, you might want to run kuavo_train/train_policy.py first. It requires the installation of the 'gym_pusht' simulation environment. Install it by running: ```bash pip install -e ".[pusht]" ``` """ import sys,os import gc from std_srvs.srv import Trigger, TriggerRequest, TriggerResponse from lerobot_patches import custom_patches from pathlib import Path from sympy import im from dataclasses import dataclass, field import hydra import gymnasium as gym import imageio import numpy import torch from tqdm import tqdm from lerobot.utils.random_utils import set_seed import datetime import time import numpy as np import json from omegaconf import DictConfig, ListConfig, OmegaConf from torchvision.transforms.functional import to_tensor from std_msgs.msg import Bool import rospy import threading import traceback from geometry_msgs.msg import PoseStamped from kuavo_deploy.config import KuavoConfig from kuavo_deploy.utils.logging_utils import setup_logger from kuavo_deploy.kuavo_service.client import PolicyClient from kuavo_deploy.utils.policy_loader import ( inject_task_prompt, load_native_policy_bundle, resolve_eval_output_dir, ) log_model = setup_logger("model") log_robot = setup_logger("robot") from kuavo_deploy.kuavo_env.KuavoSimEnv import KuavoSimEnv from kuavo_deploy.kuavo_env.KuavoRealEnv import KuavoRealEnv from kuavo_deploy.utils.ros_manager import ROSManager init_evt = threading.Event() pause_flag = threading.Event() stop_flag = threading.Event() success_evt = threading.Event() def env_init_service(req): log_robot.info(f"env_init_callback! req = {req}") init_evt.set() return TriggerResponse(success=True, message="Env init successful") def pause_callback(msg): if msg.data: pause_flag.set() else: pause_flag.clear() def stop_callback(msg): if msg.data: stop_flag.set() def env_success_callback(msg): # log_model.info("env_success_callback!") if msg.data: success_evt.set() pause_sub = rospy.Subscriber('/kuavo/pause_state', Bool, pause_callback, queue_size=10) stop_sub = rospy.Subscriber('/kuavo/stop_state', Bool, stop_callback, queue_size=10) def save_rollout_video(output_path: Path, frames: list[np.ndarray], fps: int | float) -> None: if not frames: return try: imageio.mimsave(str(output_path), frames, fps=fps, codec="libx264") except Exception as exc: log_robot.warning(f"Failed to write mp4 '{output_path.name}': {exc}. Falling back to gif.") gif_path = output_path.with_suffix(".gif") imageio.mimsave(str(gif_path), frames, fps=fps) def safe_reset_service(reset_service) -> None: """安全重置服务""" try: # 调用重置服务 response = reset_service(TriggerRequest()) if response.success: log_robot.info(f"Reset service successful: {response.message}") else: log_robot.warning(f"Reset service failed: {response.message}") except rospy.ServiceException as e: log_robot.error(f"Reset service exception: {e}") def check_control_signals(): """检查控制信号""" # 检查暂停状态 while pause_flag.is_set(): log_robot.info("🔄 机械臂运动已暂停") time.sleep(0.1) if stop_flag.is_set(): log_robot.info("🛑 机械臂运动被停止") return False # 检查是否需要停止 if stop_flag.is_set(): log_robot.info("🛑 收到停止信号，退出机械臂运动") return False return True # 正常继续 def setup_policy(pretrained_path, policy_type, device=torch.device("cuda"), task_prompt="robot manipulation"): """ Set up and load the policy model. Args: pretrained_path: Path to the checkpoint policy_type: Type of policy ('diffusion' or 'act') Returns: Loaded policy model and device """ if device.type == 'cpu': log_model.warning("Warning: Using CPU for inference, this may be slow.") time.sleep(3) if policy_type == 'client': preprocessor, postprocessor = lambda obs: obs, lambda action: action return PolicyClient(task_prompt=task_prompt), preprocessor, postprocessor, None policy, preprocessor, postprocessor, pretrained_model_dir = load_native_policy_bundle( pretrained_path=pretrained_path, device=device, strict=True, ) log_model.info(f"Model loaded from {pretrained_model_dir}") log_model.info(f"Model type: {policy.config.type}") log_model.info(f"Model n_obs_steps: {policy.config.n_obs_steps}") log_model.info(f"Model device: {device}") return policy, preprocessor, postprocessor, pretrained_model_dir def run_single_episode(config, policy, preprocessor, postprocessor, episode, output_directory): """运行单个episode""" cfg = config.inference seed = cfg.seed task_prompt = cfg.task_prompt # Initialize environment env = gym.make( config.env.env_name, max_episode_steps=cfg.max_episode_steps, config=config, ) run_single_ros_manager = ROSManager() # Setup ROS subscribers and services run_single_ros_manager.register_subscriber("/simulator/success", Bool, env_success_callback) # max_episode_steps = cfg.max_episode_steps start_service = rospy.ServiceProxy('/simulator/start', Trigger) if cfg.policy_type != 'client': log_model.info(f"policy.config.input_features: {policy.config.input_features}") log_robot.info(f"env.observation_space: {env.observation_space}") log_model.info(f"policy.config.output_features: {policy.config.output_features}") log_robot.info(f"env.action_space: {env.action_space}") # Reset the policy and environments to prepare for rollout policy.reset() observation, info = env.reset(seed=seed) if cfg.policy_type != "client": observation = inject_task_prompt(observation, task_prompt) # first_img = (observation["observation.images.head_cam_h"].squeeze().permute(1,2,0).numpy()*255).astype(np.uint8) # import cv2 # first_img = cv2.cvtColor(first_img,cv2.COLOR_RGB2BGR) # cv2.imwrite( "obs.png", first_img) # raise ValueError("stop for debug!") start_service(TriggerRequest()) # Prepare to collect every rewards and all the frames of the episode, # from initial state to final state. rewards = [] cam_keys = [k for k in observation.keys() if "images" in k or "depth" in k] frame_temp_dirs = {} for k in cam_keys: temp_dir = output_directory / f"temp_frames_{episode}_{k}" temp_dir.mkdir(parents=True, exist_ok=True) frame_temp_dirs[k] = temp_dir average_exec_time = 0 average_action_infer_time = 0 average_step_time = 0 step = 0 done = False while not done: # --- Pause support: block here if pause_flag is set --- if not check_control_signals(): log_robot.info("🛑 收到停止信号，退出机械臂运动") return 0 start_time = time.time() if cfg.policy_type != "client": observation = inject_task_prompt(observation, task_prompt) observation = preprocessor(observation) with torch.inference_mode(): action = policy.select_action(observation) log_model.info(f"Step {step}: predict action {action}") action = postprocessor(action) # print(f"action: {action}, action.shape: {action.shape}, action min: {action.min()}, action max: {action.max()}") action_infer_time = time.time() log_model.info(f"episode {episode}, step {step}, action infer time: {action_infer_time - start_time:.3f}s") average_action_infer_time += action_infer_time - start_time numpy_action = action.squeeze(0).cpu().numpy() log_model.info(f"Step {step}: Executing action {numpy_action}") observation, reward, terminated, truncated, info = env.step(numpy_action) if cfg.policy_type != "client": observation = inject_task_prompt(observation, task_prompt) exec_time = time.time() log_model.debug(f"step {step}: exec time: {exec_time - action_infer_time:.3f}s") average_exec_time += exec_time - action_infer_time rewards.append(reward) for k in cam_keys: frame_path = frame_temp_dirs[k] / f"frame_{step:04d}.png" img = (observation[k].squeeze(0).cpu().numpy().transpose(1, 2, 0) * 255).astype(np.uint8) if img.shape[-1] == 1: img = img.squeeze(-1) imageio.imwrite(str(frame_path), img) # The rollout is considered done when the success state is reached (i.e. terminated is True), # or the maximum number of iterations is reached (i.e. truncated is True) done = terminated | truncated | done done = done or success_evt.is_set() step += 1 end_time = time.time() log_model.debug(f"Step {step} time: {end_time - start_time:.3f}s") average_step_time += end_time - start_time # Get the speed of environment (i.e. its number of frames per second). fps = env.unwrapped.ros_rate log_model.info(f"average exec time: {average_exec_time / step:.3f}s") log_model.info(f"average action infer time: {average_action_infer_time / step:.3f}s") log_model.info(f"average step time: {average_step_time / step:.3f}s") log_model.info(f"average sleep time: {env.unwrapped.average_sleep_time / step:.3f}s") for cam in cam_keys: temp_dir = frame_temp_dirs[cam] frame_files = sorted(temp_dir.glob("frame_*.png")) frames = [imageio.imread(str(f)) for f in frame_files] output_path = output_directory / f"rollout_{episode}_{cam}.mp4" save_rollout_video(output_path, frames, fps) for f in frame_files: f.unlink() temp_dir.rmdir() del frames success = success_evt.is_set() env.close() run_single_ros_manager.close() del rewards del observation del env del run_single_ros_manager gc.collect() torch.cuda.empty_cache() return 1 if success else 0 # 返回是否成功 def kuavo_eval_autotest(config: KuavoConfig): """执行自动测试""" cfg = config.inference eval_episodes = cfg.eval_episodes seed = cfg.seed policy_type = cfg.policy_type # Setup paths if cfg.pretrained_path: pretrained_path = Path(cfg.pretrained_path) else: pretrained_path = Path(f"outputs/train/{cfg.task}/{cfg.method}/{cfg.timestamp}/epoch{cfg.epoch}") output_directory = resolve_eval_output_dir(pretrained_path, Path("outputs/eval")) output_directory.mkdir(parents=True, exist_ok=True) # Log evaluation results log_file_path = output_directory / "evaluation_autotest.log" with log_file_path.open("w") as log_file: log_file.write(f"Evaluation Timestamp: {datetime.datetime.now()}\n") log_file.write(f"Total Episodes: {eval_episodes}\n") # Setup policy and environment (只加载一次) set_seed(seed) device = torch.device(cfg.device) task_prompt = getattr(cfg, 'task_prompt', "robot manipulation") policy, preprocessor, postprocessor, pretrained_model_dir = setup_policy( pretrained_path, policy_type, device, task_prompt=task_prompt ) # first reset reset_service = rospy.ServiceProxy('/simulator/reset', Trigger) # Ros service init_service = rospy.Service("/simulator/init", Trigger, env_init_service) wait_times = 8 while not init_evt.is_set(): log_robot.info("Waiting for first env init...") if not check_control_signals(): log_robot.info("🛑 收到停止信号，退出机械臂运动") return time.sleep(1) wait_times -= 1 if wait_times <=0: break safe_reset_service(reset_service) init_evt.clear() success_count = 0 for episode in range(eval_episodes): while not init_evt.is_set(): log_robot.info("Waiting for env init...") if not check_control_signals(): log_robot.info("🛑 收到停止信号，退出机械臂运动") return time.sleep(1) try: result = run_single_episode(config, policy, preprocessor, postprocessor, episode, output_directory) log_robot.info(f"Episode {episode+1} completed with return code: {result}") # 重置policy状态，清理缓存 policy.reset() # 强制垃圾回收和GPU缓存清理 gc.collect() torch.cuda.empty_cache() except Exception as e: log_robot.error(f"Exception during episode {episode+1}: {e}") log_robot.error(traceback.format_exc()) result = 0 # Treat as failure safe_reset_service(reset_service) init_evt.clear() success_evt.clear() # 异常情况下也要清理内存 gc.collect() torch.cuda.empty_cache() break # 记录episode结果 episode_end_time = datetime.datetime.now().isoformat() is_success = result == 1 if is_success: success_count += 1 log_model.info(f"✅ Episode {episode+1}: Success!") else: log_model.info(f"❌ Episode {episode+1}: Failed!") with log_file_path.open("a") as log_file: log_file.write("\n") log_file.write(f"Success Count: {success_count} / Already eval episodes: {episode+1}") safe_reset_service(reset_service) init_evt.clear() success_evt.clear() # Display final statistics log_model.info("\n" + "="*50) log_model.info(f"🎯 Evaluation completed!") log_model.info(f"📊 Success count: {success_count}/{eval_episodes}") log_model.info(f"📈 Success rate: {success_count / eval_episodes:.2%}") log_model.info(f"📁 Videos and logs saved to: {output_directory}") log_model.info("="*50) init_service.shutdown() pause_sub.unregister() stop_sub.unregister()