Mujoco_WASM/training_server/mobile_manipulator/evaluation.py

"""Seeded held-out rollouts; successful export is never a task-success metric."""

from .env import MobileManipulatorEnv


def evaluate_policy(agent, package, params, reset_options, seed):
    env = MobileManipulatorEnv(
        package,
        stage=params["stage"],
        reset_options=reset_options,
        position_jitter=params["positionJitter"],
    )
    episodes = []
    try:
        for i in range(params["evaluationEpisodes"]):
            observation, _ = env.reset(seed=seed + i)
            done = False
            reward_sum = 0.0
            while not done:
                action, _ = agent.predict(observation, deterministic=True)
                observation, reward, terminated, truncated, info = env.step(action)
                done = terminated or truncated
                reward_sum += reward
            episodes.append(
                {
                    "seed": seed + i,
                    "success": info["is_success"],
                    "safetyStop": info["safety_stop"],
                    "steps": env.kernel.steps,
                    "maxJointVelocity": info["max_joint_velocity"],
                    "navigationDistance": info["navigation_distance"],
                    "reward": reward_sum,
                }
            )
            print(
                f"Evaluation episode {i + 1}/{params['evaluationEpisodes']}: "
                f"success={info['is_success']} safety={info['safety_stop'] or 'none'} "
                f"max_joint_velocity={info['max_joint_velocity']:.5f}",
                flush=True,
            )
    finally:
        env.close()
    return {
        "episodes": len(episodes),
        "successRate": sum(x["success"] for x in episodes) / len(episodes),
        "safetyStops": sum(bool(x["safetyStop"]) for x in episodes),
        "maxJointVelocity": max(x["maxJointVelocity"] for x in episodes),
        "meanNavigationDistance": sum(x["navigationDistance"] for x in episodes) / len(episodes),
        "seed": seed,
        "positionJitter": params["positionJitter"],
        "rollouts": episodes,
    }