bdfc77147d
Signed-off-by: Meng <meng_chen99@163.com> Signed-off-by: Yangshen Deng <yangshen.d@outlook.com> Signed-off-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: Meng <meng_chen99@163.com> Co-authored-by: Suli Wang <50828270+wsl2000@users.noreply.github.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
87 lines
3.0 KiB
Python
87 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
os.environ.setdefault("MUJOCO_GL", "egl")
|
|
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
|
|
|
|
# Import base configs at module top level so the subclasses below are pickle-
|
|
# resolvable (worker processes import this module fresh via __main__).
|
|
from molmo_spaces.configs.policy_configs_baselines import ( # noqa: E402
|
|
DreamZeroPolicyConfig,
|
|
)
|
|
from molmo_spaces.evaluation.configs.evaluation_configs import ( # noqa: E402
|
|
DreamZeroPolicyEvalConfig,
|
|
)
|
|
|
|
|
|
# We only need to change the backend host and port to the vllm-host!
|
|
class DreamZeroVllmOmniPolicyConfig(DreamZeroPolicyConfig):
|
|
remote_config: dict = dict(host="127.0.0.1", port=8000)
|
|
|
|
|
|
class DreamZeroVllmOmniEvalConfig(DreamZeroPolicyEvalConfig):
|
|
policy_config: DreamZeroVllmOmniPolicyConfig = DreamZeroVllmOmniPolicyConfig()
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="127.0.0.1")
|
|
parser.add_argument("--port", type=int, default=8000)
|
|
parser.add_argument(
|
|
"--benchmark_dir",
|
|
required=True,
|
|
help=(
|
|
"Path to a MolmoSpaces benchmark directory, for example "
|
|
"$MOLMOSPACES_BENCHMARK_DIR/20260327/ithor/FrankaCloseHardBench/"
|
|
"FrankaCloseHardBench_20260206_json_benchmark"
|
|
),
|
|
)
|
|
parser.add_argument("--max_episodes", type=int, default=1)
|
|
parser.add_argument("--task_horizon_steps", type=int, default=80)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
required=True,
|
|
help="Directory to write evaluation outputs (created if missing).",
|
|
)
|
|
parser.add_argument("--episode_idx", type=int, default=None)
|
|
args = parser.parse_args()
|
|
|
|
policy_config = DreamZeroVllmOmniPolicyConfig(remote_config=dict(host=args.host, port=args.port))
|
|
DreamZeroVllmOmniPolicyConfig.model_fields["remote_config"].default = policy_config.remote_config
|
|
DreamZeroVllmOmniEvalConfig.model_fields["policy_config"].default = policy_config
|
|
|
|
# Import after env vars are set so MuJoCo picks EGL.
|
|
from molmo_spaces.evaluation import run_evaluation
|
|
|
|
cfg_cls = DreamZeroVllmOmniEvalConfig
|
|
|
|
output_dir = args.output_dir
|
|
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"[eval] benchmark_dir={args.benchmark_dir}")
|
|
print(f"[eval] max_episodes={args.max_episodes} task_horizon_steps={args.task_horizon_steps}")
|
|
print(f"[eval] remote policy: ws://{args.host}:{args.port}/v1/realtime/robot/openpi")
|
|
|
|
results = run_evaluation(
|
|
eval_config_cls=cfg_cls,
|
|
benchmark_dir=Path(args.benchmark_dir),
|
|
max_episodes=args.max_episodes,
|
|
task_horizon_steps=args.task_horizon_steps,
|
|
num_workers=1,
|
|
use_wandb=False,
|
|
output_dir=output_dir,
|
|
episode_idx=args.episode_idx,
|
|
)
|
|
|
|
print(f"[eval] success={results.success_count}/{results.total_count} ({results.success_rate:.1%})")
|
|
print(f"[eval] output_dir={results.output_dir}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|