eval_traj.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import json
import os

from reward_function import get_config_model_name, score_trajectory


def main() -> None:
    parser = argparse.ArgumentParser(description="Single trajectory evaluator")
    parser.add_argument("--traj-dir", type=str, required=True, help="Path to single trajectory directory")
    parser.add_argument(
        "--config",
        type=str,
        default="config.yaml",
        help="Path to config yaml. Defaults to ./config.yaml (current working directory).",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="",
        help="Directory for evaluation artifacts. Defaults to ./multi_agent_eval_outputs/<model_name>/<task_name>/",
    )
    args = parser.parse_args()

    config_path = args.config
    run_name = get_config_model_name(config_path)
    resolved_output_dir = args.output_dir.strip()
    if not resolved_output_dir:
        task_name = os.path.basename(os.path.abspath(args.traj_dir).rstrip("/"))
        resolved_output_dir = os.path.join(
            os.getcwd(),
            "multi_agent_eval_outputs",
            run_name,
            task_name,
        )

    result = score_trajectory(
        traj_dir=args.traj_dir,
        agent_name=None,
        config_path=config_path,
        output_dir=resolved_output_dir,
    )
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()