test.py

import os
import collections
import math
import collections
import dataclasses
import logging
import tyro

import numpy as np
import pickle as pkl

from tqdm import tqdm, trange
from pathlib import Path
from libero.libero import benchmark
from libero.libero import get_libero_path
from libero.libero.envs import OffScreenRenderEnv
from openpi_client import image_tools
from openpi_client import websocket_client_policy as _websocket_client_policy

def _quat2axisangle(quat):
    """
    Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
    """
    # 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 collect_pi0fast_embedding(metadata):
    # using best ablation settings from SAFE paper
    pre_logits = metadata["pre_logits"]
    return np.mean(pre_logits, axis=0)

def pi0fast_embedding_eval(embeddings):
    norms = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
    cos_sim = norms @ norms.T
    n = len(embeddings)
    mean_cosine = (np.sum(cos_sim) - n) / (n * (n - 1))
    print(mean_cosine)

def evaluate_embedding(embedding):
    # USE AURORA
    pass

def simulate_libero_base(experiment_cfg):
    np.random.seed(experiment_cfg["seed"])

    # initialize model
    client = _websocket_client_policy.WebsocketClientPolicy(experiment_cfg["host"], experiment_cfg["port"])

    # initialize libero env
    benchmark_dict = benchmark.get_benchmark_dict()
    task_suite = benchmark_dict["libero_spatial"]()

    task_id = -1
    for i in range(10):
        task = task_suite.get_task(i)
        if task.language == "pick up the black bowl next to the ramekin and place it on the plate":
            task_id = i
            break

    initial_states = task_suite.get_task_init_states(task_id)
    task = task_suite.get_task(task_id)
    task_description = task.language

    task_bddl = Path(get_libero_path("bddl_files")) / task.problem_folder / task.bddl_file
    env_args = {
        "bddl_file_name": task_bddl, 
        "camera_heights": 256, 
        "camera_widths": 256
    }

    env = OffScreenRenderEnv(**env_args)
    env.seed = experiment_cfg["seed"]

    success_rate = 0

    embeddings = []

    for episode_idx in trange(experiment_cfg["num_trials_per_task"]):
        env.reset()
        action_plan = collections.deque()

        if initial_states is None:
            obs = env.env._get_observations()
        else:
            obs = env.set_init_state(initial_states[episode_idx])

        success = False
        t = 0
        while t < experiment_cfg["max_steps"] + experiment_cfg["num_steps_wait"]:
            try:
                if t < experiment_cfg["num_steps_wait"]:
                    obs, reward, done, info = env.step([0.0] * 6 + [-1.0])
                    t += 1
                    continue

                # Get preprocessed image
                # IMPORTANT: rotate 180 degrees to match train preprocessing
                img = np.ascontiguousarray(obs["agentview_image"][::-1, ::-1])
                wrist_img = np.ascontiguousarray(obs["robot0_eye_in_hand_image"][::-1, ::-1])

                if not action_plan:
                    # Finished executing previous action chunk -- compute new chunk
                    # Prepare observations dict
                    element = {
                        "observation/image": img,
                        "observation/wrist_image": wrist_img,
                        "observation/state": np.concatenate(
                            (
                                obs["robot0_eef_pos"],
                                _quat2axisangle(obs["robot0_eef_quat"]),
                                obs["robot0_gripper_qpos"],
                            )
                        ),
                        "prompt": env.language_instruction,
                        "run/env_id": experiment_cfg["env_id"]
                    }

                    # Query model to get action
                    if t == experiment_cfg["num_steps_wait"]:
                        action_embed = client.infer(element)
                        initial_embeddings = collect_pi0fast_embedding(action_embed)

                        embeddings.append(initial_embeddings)

                        action_chunk = action_embed["actions"]
                    else:
                        action_chunk = client.infer(element)["actions"]
                    
                    action_chunk = np.squeeze(action_chunk)

                    assert (
                        len(action_chunk) >= experiment_cfg["replan_steps"]
                    ), f"We want to replan every {experiment_cfg['replan_steps']} steps, but policy only predicts {len(action_chunk)} steps."
                    action_plan.extend(action_chunk[: experiment_cfg["replan_steps"]])

                action = action_plan.popleft()

                # Execute action in environment
                obs, reward, done, info = env.step(action.tolist())
                if done:
                    success_rate += 1 / experiment_cfg["num_trials_per_task"]
                    success = True
                    break
                t += 1
            except Exception as e:
                print(e)
                break
    
    pi0fast_embedding_eval(embeddings)

    return success_rate

experiment_cfg = {
    "host": "0.0.0.0",
    "port": 8000,
    "resize_size": 224,
    "replan_steps": 1,

    "task_suite_name": "libero_spatial",
    "num_steps_wait": 10,
    "num_trials_per_task": 5,
    "max_steps": 220,
    "env_id": 1,

    "save_name": "pi0fast-test",
    "seed": 42
}

if __name__ == "__main__":
    simulate_libero_base(experiment_cfg)