eval_libero.py

import os
import sys
import numpy as np
import torch
import tqdm
import imageio
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Tuple, Optional
import os, sys, tqdm, numpy as np, torch
import tensorflow as tf
import math
# 只用 GPU4
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
sys.path.append("/data/wuyuzhou-20250917/openpi/third_party/libero")
from libero.libero import benchmark
from libero.libero.envs import OffScreenRenderEnv
from libero.libero import get_libero_path
from libero.libero.envs import OffScreenRenderEnv
# 添加当前目录到系统路径
#sys.path.append("../..")

# 注意：您需要实现_config和policy_config模块
# 假设这些模块在您的环境中可用
from openpi.training import config as _config
from openpi.policies import policy_config

@dataclass
class TestConfig:
    """测试配置参数"""
    # 任务套件名称
    task_suite_name: str = "libero_spatial"
    
    # 每个任务的试验次数
    num_trials_per_task: int = 10
    
    # 等待物体稳定的步数
    num_steps_wait: int = 10
    
    # 日志目录
    local_log_dir: str = "./test_logs"
    
    # 视频保存目录
    video_dir: str = "./test_videos"
    
    # 是否保存视频
    save_videos: bool = True
    
    # 随机种子
    seed: int = 42
    
    # 最大步数（根据任务类型设置）
    max_steps: int = 220  # libero_spatial的默认最大步数
    
    # 策略配置
    policy_name: str = "pi05_libero"
    checkpoint_dir: str = "/data/wuyuzhou-20250917/openpi/checkpoints/pi05_libero/test_single/30000"
    device: str = "cuda:0"
    
    # 任务特定配置
    task_configs: Dict[str, Dict] = field(default_factory=lambda: {
        "libero_spatial": {"max_steps": 220},
        "libero_object": {"max_steps": 280},
        "libero_goal": {"max_steps": 300},
        "libero_10": {"max_steps": 520},
        "libero_90": {"max_steps": 400},
    })

def get_image_resize_size(cfg):
    """
    Gets image resize size for a model class.
    If `resize_size` is an int, then the resized image will be a square.
    Else, the image will be a rectangle.
    """
    # if cfg.model_family == "openvla":
    #     resize_size = 224
    # else:
    #     raise ValueError("Unexpected `model_family` found in config.")
    return 224

def get_libero_env(task,  resolution=256):
    """Initializes and returns the LIBERO environment, along with the task description."""
    task_description = task.language
    task_bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file)
    env_args = {"bddl_file_name": task_bddl_file, "camera_heights": resolution, "camera_widths": resolution}
    env = OffScreenRenderEnv(**env_args)
    env.seed(0)  # IMPORTANT: seed seems to affect object positions even when using fixed initial state
    return env, task_description

def load_policy():
    config = _config.get_config("pi05_libero")
    checkpoint_dir = "/data/wuyuzhou-20250917/openpi/checkpoints/pi05_libero/test_single/30000"
    policy = policy_config.create_trained_policy(config, checkpoint_dir, pytorch_device="cuda:0")
    return policy


def format_observation(obs: dict, task_description: str) -> dict:
    """格式化观察数据以匹配策略输入"""
    return {
        "observation/image": np.asarray(obs["agentview_image"], dtype=np.float32),
        "observation/wrist_image": np.asarray(obs["robot0_eye_in_hand_image"], dtype=np.float32),
        "observation/state": np.concatenate([
            obs["robot0_eef_pos"],
            obs["robot0_eef_quat"],
            obs["robot0_gripper_qpos"],
        ]).astype(np.float32),
        "prompt": task_description,
    }


def get_libero_dummy_action(model_family: str):
    """Get dummy/no-op action, used to roll out the simulation while the robot does nothing."""
    return [0, 0, 0, 0, 0, 0, -1]


def resize_image(img, resize_size):
    """
    Takes numpy array corresponding to a single image and returns resized image as numpy array.

    NOTE (Moo Jin): To make input images in distribution with respect to the inputs seen at training time, we follow
                    the same resizing scheme used in the Octo dataloader, which OpenVLA uses for training.
    """
    assert isinstance(resize_size, tuple)
    # Resize to image size expected by model
    img = tf.image.encode_jpeg(img)  # Encode as JPEG, as done in RLDS dataset builder
    img = tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8)  # Immediately decode back
    img = tf.image.resize(img, resize_size, method="lanczos3", antialias=True)
    img = tf.cast(tf.clip_by_value(tf.round(img), 0, 255), tf.uint8)
    img = img.numpy()
    return img

def get_libero_image(obs, resize_size):
    """Extracts image from observations and preprocesses it."""
    assert isinstance(resize_size, int) or isinstance(resize_size, tuple)
    if isinstance(resize_size, int):
        resize_size = (resize_size, resize_size)
    img = obs["agentview_image"]
    img = img[::-1, ::-1]  # IMPORTANT: rotate 180 degrees to match train preprocessing
    img = resize_image(img, resize_size)
    return img

def quat2axisangle(quat):
    """
    Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55

    Converts quaternion to axis-angle format.
    Returns a unit vector direction scaled by its angle in radians.

    Args:
        quat (np.array): (x,y,z,w) vec4 float angles

    Returns:
        np.array: (ax,ay,az) axis-angle exponential coordinates
    """
    # clip quaternion
    if quat[3] > 1.0:
        quat[3] = 1.0
    elif quat[3] < -1.0:
        quat[3] = -1.0

    den = np.sqrt(1.0 - quat[3] * quat[3])
    if math.isclose(den, 0.0):
        # This is (close to) a zero degree rotation, immediately return
        return np.zeros(3)

    return (quat[:3] * 2.0 * math.acos(quat[3])) / den


def run_libero_test(cfg: TestConfig):
    """运行LIBERO测试套件"""
    # 创建日志目录
    os.makedirs(cfg.local_log_dir, exist_ok=True)
    log_filepath = os.path.join(cfg.local_log_dir, f"test_log_{cfg.task_suite_name}.txt")
    log_file = open(log_filepath, "w")
    
    # 创建视频目录
    if cfg.save_videos:
        os.makedirs(cfg.video_dir, exist_ok=True)
    
    # 加载策略
    print("Loading policy...")
    policy = load_policy()
    print("Policy loaded successfully.")
    log_file.write(f"Policy loaded from: {cfg.checkpoint_dir}\n")
    
    # 设置随机种子
    np.random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)
    
    # 根据任务类型设置最大步数
    if cfg.task_suite_name in cfg.task_configs:
        cfg.max_steps = cfg.task_configs[cfg.task_suite_name]["max_steps"]
    
    # 初始化LIBERO任务套件
    benchmark_dict = benchmark.get_benchmark_dict()
    task_suite = benchmark_dict[cfg.task_suite_name]()
    num_tasks_in_suite = task_suite.n_tasks
    resize_size = get_image_resize_size(cfg)
    print(f"Starting LIBERO {cfg.task_suite_name} test with {num_tasks_in_suite} tasks")
    log_file.write(f"Starting LIBERO {cfg.task_suite_name} test with {num_tasks_in_suite} tasks\n")
    
    # 初始化统计
    total_episodes, total_successes = 0, 0
    task_results = {}
    
    # 遍历所有任务
    for task_id in tqdm.tqdm(range(num_tasks_in_suite), desc="Tasks"):
        # 获取任务
        task = task_suite.get_task(task_id)
        task_description = task.language
        print(f"\nTesting task: {task_description}")
        log_file.write(f"\nTesting task: {task_description}\n")
        
        # 获取默认初始状态
        initial_states = task_suite.get_task_init_states(task_id)
        env, task_description = get_libero_env(task, resolution=256)
        
        # 任务统计
        task_episodes, task_successes = 0, 0
        
        # 遍历每个试验
        for episode_idx in tqdm.tqdm(range(cfg.num_trials_per_task), desc="Trials", leave=False):
            # 初始化环境
            #env = OffScreenRenderEnv(task, render_mode="rgb_array")
            
            # 重置环境
            env.reset()
            obs=env.set_init_state(initial_states[episode_idx])
            replay_images = []
            # 初始化视频记录
            video_frames = []
            
            # # 等待物体稳定
            # for _ in range(cfg.num_steps_wait):
            #     obs, _, _, _ = env.step(env.get_null_action())
            #     if cfg.save_videos:
            #         video_frames.append(obs["agentview_image"])
            
            # 主循环
            done = False
            step_count = 0
            
            while not done and step_count < cfg.max_steps:
                # 获取观察
                img = get_libero_image(obs, resize_size)

                # Save preprocessed image for replay video
                replay_images.append(img)

                    # Prepare observations dict
                    # Note: OpenVLA does not take proprio state as input
                observation = {
                        "full_image": img,
                        "state": np.concatenate(
                            (obs["robot0_eef_pos"], quat2axisangle(obs["robot0_eef_quat"]), obs["robot0_gripper_qpos"])
                        ),
                    }
                
                # 格式化观察
                # formatted_obs = format_observation(obs, task_description)
                
                # 使用策略推断动作
                with torch.inference_mode():
                    action = policy.infer(observation)["actions"]
                
                # 执行动作
                obs, reward, done, info = env.step(action)
                
                # 记录视频帧
                if cfg.save_videos:
                    video_frames.append(obs["agentview_image"])
                
                step_count += 1
            
            # 关闭环境
            # env.close()
            
            # 更新统计
            task_episodes += 1
            total_episodes += 1
            if done:
                task_successes += 1
                total_successes += 1
            
            # 保存视频
            if cfg.save_videos and video_frames:
                video_path = os.path.join(
                    cfg.video_dir, 
                    f"{cfg.task_suite_name}_task_{task_id}_trial_{episode_idx}.mp4"
                )
                imageio.mimsave(video_path, video_frames, fps=30)
        
        # 记录任务结果
        task_success_rate = task_successes / task_episodes if task_episodes > 0 else 0.0
        task_results[task_description] = {
            "success_rate": task_success_rate,
            "successes": task_successes,
            "episodes": task_episodes
        }
        
        print(f"Task '{task_description}' success rate: {task_success_rate:.2f} ({task_successes}/{task_episodes})")
        log_file.write(f"Task '{task_description}' success rate: {task_success_rate:.2f} ({task_successes}/{task_episodes})\n")
    
    # 计算总体成功率
    overall_success_rate = total_successes / total_episodes if total_episodes > 0 else 0.0
    
    # 打印和记录测试摘要
    print("\n=== Test Summary ===")
    print(f"Task suite: {cfg.task_suite_name}")
    print(f"Total episodes: {total_episodes}")
    print(f"Total successes: {total_successes}")
    print(f"Overall success rate: {overall_success_rate:.2f}")
    
    log_file.write("\n=== Test Summary ===\n")
    log_file.write(f"Task suite: {cfg.task_suite_name}\n")
    log_file.write(f"Total episodes: {total_episodes}\n")
    log_file.write(f"Total successes: {total_successes}\n")
    log_file.write(f"Overall success rate: {overall_success_rate:.2f}\n")
    
    # 打印详细任务结果
    print("\nDetailed Task Results:")
    log_file.write("\nDetailed Task Results:\n")
    for task, result in task_results.items():
        print(f"  - {task}: {result['success_rate']:.2f} ({result['successes']}/{result['episodes']})")
        log_file.write(f"  - {task}: {result['success_rate']:.2f} ({result['successes']}/{result['episodes']})\n")
    
    # 关闭日志文件
    log_file.close()
    
    return {
        "task_suite": cfg.task_suite_name,
        "total_episodes": total_episodes,
        "total_successes": total_successes,
        "overall_success_rate": overall_success_rate,
        "task_results": task_results
    }


if __name__ == "__main__":
    # 配置测试参数
    cfg = TestConfig(
        task_suite_name="libero_spatial",
        num_trials_per_task=5,
        save_videos=True,
        video_dir="./libero_spatial_test_videos"
    )
    
    # 运行测试
    print(f"Starting LIBERO {cfg.task_suite_name} test...")
    results = run_libero_test(cfg)
    
    # 打印结果摘要
    print("\n=== Final Test Summary ===")
    print(f"Task suite: {results['task_suite']}")
    print(f"Total episodes: {results['total_episodes']}")
    print(f"Total successes: {results['total_successes']}")
    print(f"Overall success rate: {results['overall_success_rate']:.2f}")