Robot benchmark

from hud import Environment
from hud.environment.robot import RobotEndpoint
from config import build_contract            # the env's own contract helper
from libero_sim_bridge import LiberoSimBridge

CONTRACT = build_contract({"use_delta": True})  # the env's self-describing obs/action schema

env = Environment(name="libero")
endpoint = RobotEndpoint(LiberoSimBridge(use_delta=True))  # drive the bridge through the endpoint

@env.initialize
async def _up():
    await endpoint.start()
    env.add_capability(await endpoint.capability(contract=CONTRACT))

@env.shutdown
async def _down():
    await endpoint.stop()

@env.template(id="libero_spatial")
async def libero_spatial(libero_task_id: int, init_state_id: int = 0):
    prompt = await endpoint.reset(task_suite="libero_spatial",
                                  task_id=libero_task_id, init_state_id=init_state_id)
    yield {"prompt": prompt}
    yield await endpoint.result()

import asyncio
import torch
from lerobot.policies.factory import make_pre_post_processors
from lerobot.policies.pi05.modeling_pi05 import PI05Policy

from hud.agents.robot.adapter import LeRobotAdapter
from hud.agents.robot.agent import RobotAgent
from hud.agents.robot.model import LeRobotModel
from hud.eval import DockerRuntime, Task, Taskset

CHECKPOINT = "lerobot/pi05_libero_finetuned"

class PI05Agent(RobotAgent):
    max_steps = 400

    def __init__(self):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        policy = PI05Policy.from_pretrained(CHECKPOINT).to(device).eval()
        pre, post = make_pre_post_processors(
            policy.config, CHECKPOINT,
            preprocessor_overrides={"device_processor": {"device": device}},
        )
        self.model = LeRobotModel(policy, pre, post)
        self.adapter = LeRobotAdapter(model_image_keys=list(policy.config.image_features))

TASKS = [
    Task(env="libero", id="libero_spatial", args={"libero_task_id": t, "init_state_id": 0})
    for t in range(3)
]

async def main():
    job = await Taskset("libero-demo", TASKS).run(
        PI05Agent(), runtime=DockerRuntime("hud-libero-env"), max_concurrent=1,
    )
    rewards = [run.reward for run in job.runs]
    print(f"success_rate={sum(rewards) / len(rewards):.2f}")

asyncio.run(main())