code/kuavo_deploy/src/eval/sim_auto_test.py

# 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 <https://www.gnu.org/licenses/>.
#
# ---
#
# 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()