run_eval.py

import argparse
import importlib
import json
import os
from pathlib import Path

from dotenv import load_dotenv
from omegaconf import OmegaConf
from habitat.config import read_write
from habitat.config.default_structured_configs import (
    AgentConfig,
    EnvironmentConfig,
    HabitatSimV0Config,
    LookDownActionConfig,
    LookUpActionConfig,
    MoveForwardActionConfig,
    RendererConfig,
    SimulatorConfig,
    StopActionConfig,
    TaskConfig,
    TurnLeftActionConfig,
    TurnRightActionConfig,
)

try:
    from habitat.core.env import Env
except ImportError:
    from habitat.core.env import Env

from dataset import load_custom_dataset
from evaluator import evaluate
from policy import HumanPolicy, ModelApiPolicy, RandomPolicy


DEFAULT_DATASET_DIR = Path("/home/wenbofu/habitat_benchmark/dataset_semantic")
DEFAULT_HM3D_SCENE_ROOT = Path("/home/wenbofu/datasets/hm3d/versioned_data/hm3d-0.2/hm3d/val")


def load_habitat_config(config_path: str):
    config = OmegaConf.load(config_path)
    if "habitat" not in config:
        raise ValueError("Habitat config must include a top-level 'habitat' section.")

    habitat_cfg = config.habitat
    with read_write(habitat_cfg.simulator):
        habitat_cfg.simulator = OmegaConf.merge(
            OmegaConf.structured(SimulatorConfig),
            habitat_cfg.simulator,
        )
        if not habitat_cfg.simulator.get("agents_order") and len(habitat_cfg.simulator.agents) == 1:
            habitat_cfg.simulator.agents_order = list(habitat_cfg.simulator.agents.keys())
        for agent_name, agent_cfg in habitat_cfg.simulator.agents.items():
            habitat_cfg.simulator.agents[agent_name] = OmegaConf.merge(
                OmegaConf.structured(AgentConfig),
                agent_cfg,
            )
        habitat_cfg.simulator.habitat_sim_v0 = OmegaConf.merge(
            OmegaConf.structured(HabitatSimV0Config),
            habitat_cfg.simulator.get("habitat_sim_v0", {}),
        )
        habitat_cfg.simulator.renderer = OmegaConf.merge(
            OmegaConf.structured(RendererConfig),
            habitat_cfg.simulator.get("renderer", {}),
        )

    with read_write(habitat_cfg):
        habitat_cfg.environment = OmegaConf.merge(
            OmegaConf.structured(EnvironmentConfig),
            habitat_cfg.get("environment", {}),
        )
        task_cfg = habitat_cfg.get("task", {})
        if not task_cfg.get("actions"):
            look_up = OmegaConf.structured(LookUpActionConfig)
            look_down = OmegaConf.structured(LookDownActionConfig)
            look_up.tilt_angle = 30
            look_down.tilt_angle = 30
            task_cfg.actions = {
                "stop": OmegaConf.structured(StopActionConfig),
                "move_forward": OmegaConf.structured(MoveForwardActionConfig),
                "turn_left": OmegaConf.structured(TurnLeftActionConfig),
                "turn_right": OmegaConf.structured(TurnRightActionConfig),
                "look_up": look_up,
                "look_down": look_down,
            }
        task_base = OmegaConf.structured(TaskConfig)
        OmegaConf.set_struct(task_base, False)
        habitat_cfg.task = OmegaConf.merge(task_base, task_cfg)
        if habitat_cfg.get("seed") is None:
            habitat_cfg.seed = 0

    return config


def parse_args():
    parser = argparse.ArgumentParser(description="Run custom HM3D multi-object navigation evaluation.")
    parser.add_argument(
        "--config", required=True, help="Path to a Habitat config file for HM3D and simulator setup."
    )
    parser.add_argument(
        "--limit", type=int, default=None, help="Limit to the first N scenes (default: all scenes)."
    )
    parser.add_argument(
        "--episode-limit",
        type=int,
        default=None,
        help="Limit to the first N episodes per scene (default: all episodes).",
    )
    parser.add_argument(
        "--dataset-dir",
        default=str(DEFAULT_DATASET_DIR),
        help="Directory containing per-scene custom episode folders.",
    )
    parser.add_argument(
        "--scene-root",
        default=str(DEFAULT_HM3D_SCENE_ROOT),
        help="HM3D split directory containing numbered scene folders.",
    )
    parser.add_argument(
        "--max-steps", type=int, default=500, help="Maximum number of steps per episode."
    )
    parser.add_argument(
        "--threshold", type=float, default=1.0, help="Proximity threshold for object success in meters."
    )
    parser.add_argument(
        "--seed", type=int, default=0, help="Random seed for the policy."
    )
    parser.add_argument(
        "--agent",
        default="auto",
        help=(
            "Agent to evaluate: auto, random, api, human, or an import path like "
            "my_package.my_module:MyAgent. Custom agents should expose act(observation)."
        ),
    )
    parser.add_argument(
        "--legacy-agent-api",
        action="store_true",
        help=(
            "Use the old custom-agent call style act(rgb=..., depth=..., prompt=...). "
            "By default, imported custom agents receive benchmark_api.Observation."
        ),
    )
    parser.add_argument(
        "--instruction-policy",
        choices=("first", "random", "all"),
        default="first",
        help=(
            "How to use multiple natural_language descriptions for one query: "
            "first uses the first one, random samples one, all concatenates all descriptions."
        ),
    )
    parser.add_argument(
        "--goal-type",
        choices=("description", "object"),
        default="description",
        help=(
            "Evaluation goal modality. description uses natural-language instructions. "
            "object uses only simple dataset-category queries where query_program is category == value."
        ),
    )
    parser.add_argument(
        "--action-protocol",
        choices=("goat", "legacy", "compat"),
        default="goat",
        help=(
            "Action semantics for custom agents. goat: 0=finish, 6=target_found, "
            "9=legacy finish alias. legacy/compat preserve old many/all 0=confirm behavior."
        ),
    )
    parser.add_argument(
        "--space-coverage-samples",
        type=int,
        default=200,
        help=(
            "Number of navmesh points sampled per episode for the approximate "
            "OER-space/FTR unexplored-space oracle. Use 0 to disable OER-space."
        ),
    )
    parser.add_argument(
        "--space-coverage-radius",
        type=float,
        default=1.5,
        help="Euclidean radius in meters used to mark sampled navmesh points as explored.",
    )
    parser.add_argument(
        "--log-episodes",
        action="store_true",
        help="Log per-episode success, coverage, and reached object IDs.",
    )
    parser.add_argument(
        "--log-actions",
        action="store_true",
        help="Log policy-selected actions for debugging API evaluation flow.",
    )
    parser.add_argument(
        "--log-prompts",
        action="store_true",
        help="Log the exact task prompt sent to the agent at the start of each episode.",
    )
    parser.add_argument(
        "--log-positions",
        action="store_true",
        help="Log the agent start position and position after each executed action.",
    )
    parser.add_argument(
        "--log-action-limit",
        type=int,
        default=100,
        help="Maximum number of policy actions to log when --log-actions is set. Use -1 for no limit.",
    )
    parser.add_argument(
        "--save-api-images-dir",
        default=None,
        help="Directory to save the RGB/depth images sent to the evaluation API. Disabled by default.",
    )
    parser.add_argument(
        "--render-live",
        action="store_true",
        help="Show a live first-person RGB window during evaluation. Press q or Esc to stop.",
    )
    parser.add_argument(
        "--render-delay-ms",
        type=int,
        default=30,
        help="Delay between live rendered frames in milliseconds.",
    )
    return parser.parse_args()


def _load_import_path(import_path: str):
    if ":" in import_path:
        module_name, object_name = import_path.split(":", 1)
    else:
        module_name, object_name = import_path.rsplit(".", 1)
    module = importlib.import_module(module_name)
    return getattr(module, object_name)


def _instantiate_custom_agent(agent_cls, action_space, args):
    attempts = (
        lambda: agent_cls(action_space=action_space, seed=args.seed),
        lambda: agent_cls(action_space=action_space),
        lambda: agent_cls(args=args),
        lambda: agent_cls(),
    )
    last_error = None
    for attempt in attempts:
        try:
            return attempt()
        except TypeError as exc:
            last_error = exc
    raise TypeError(f"Could not instantiate custom agent {agent_cls!r}: {last_error}") from last_error


def make_agent(agent_name: str, action_space, args):
    agent_name = str(agent_name or "auto")
    api_key = os.environ.get("EVAL_MODEL_API_KEY")
    api_url = os.environ.get("EVAL_MODEL_API_URL")
    model_name = os.environ.get("EVAL_MODEL_NAME")

    if agent_name == "auto":
        agent_name = "api" if api_key and api_url and model_name else "random"

    if agent_name == "random":
        agent = RandomPolicy(action_space, seed=args.seed)
        print("Using RandomPolicy.")
    elif agent_name == "human":
        agent = HumanPolicy(action_space, seed=args.seed)
        print("Using HumanPolicy. Enter actions in the terminal.")
    elif agent_name == "api":
        if not (api_key and api_url and model_name):
            raise ValueError(
                "--agent api requires EVAL_MODEL_API_KEY, EVAL_MODEL_API_URL, and EVAL_MODEL_NAME."
            )
        agent = ModelApiPolicy(
            action_space,
            api_key=api_key,
            api_url=api_url,
            model_name=model_name,
            seed=args.seed,
        )
        print("Using ModelApiPolicy via environment API configuration.")
    else:
        agent_cls = _load_import_path(agent_name)
        agent = _instantiate_custom_agent(agent_cls, action_space, args)
        if not args.legacy_agent_api:
            setattr(agent, "_benchmark_observation_api", True)
        print(f"Using custom agent: {agent_name}")

    for name, value in (
        ("log_actions", args.log_actions),
        ("log_action_limit", None if args.log_action_limit < 0 else args.log_action_limit),
        ("save_api_images_dir", args.save_api_images_dir),
    ):
        try:
            setattr(agent, name, value)
        except AttributeError:
            pass
    return agent


def resolve_hm3d_scene_path(scene_id: str, scene_root: Path) -> Path:
    scene_id = str(scene_id).strip()
    if not scene_id:
        raise ValueError("Episode is missing scene_id.")

    scene_path = Path(scene_id)
    if scene_path.suffix == ".glb":
        if scene_path.exists():
            return scene_path
        raise FileNotFoundError(f"Scene file does not exist: {scene_path}")

    matches = sorted(scene_root.glob(f"*-{scene_id}/{scene_id}.basis.glb"))
    if not matches:
        matches = sorted(scene_root.glob(f"{scene_id}/{scene_id}.basis.glb"))
    if len(matches) != 1:
        match_text = "no matches" if not matches else ", ".join(str(path) for path in matches)
        raise FileNotFoundError(f"Could not resolve HM3D scene '{scene_id}' under {scene_root}: {match_text}")
    return matches[0]


def attach_resolved_scene_paths(dataset, scene_root: Path) -> None:
    scene_cache = {}
    for episode in dataset.episodes:
        if episode.scene_id not in scene_cache:
            scene_cache[episode.scene_id] = str(resolve_hm3d_scene_path(episode.scene_id, scene_root))
        episode.scene_id = scene_cache[episode.scene_id]


def main():
    load_dotenv(Path(__file__).with_name(".env"))
    args = parse_args()
    config = load_habitat_config(args.config)

    dataset_dir = Path(args.dataset_dir)
    scene_root = Path(args.scene_root)
    scene_dirs = sorted([d for d in dataset_dir.iterdir() if d.is_dir()])

    if args.limit is not None:
        scene_dirs = scene_dirs[:args.limit]

    all_episode_results = []

    env = Env(config)
    policy = make_agent(args.agent, env.action_space, args)

    for scene_dir in scene_dirs:
        episodes_path = scene_dir / "episodes.json"
        queries_path = scene_dir / "queries.json"
        goals_path = scene_dir / "goals.json"

        if not episodes_path.exists():
            print(f"Warning: {episodes_path} not found, skipping {scene_dir.name}")
            continue

        dataset = load_custom_dataset(
            str(episodes_path),
            str(queries_path) if queries_path.exists() else None,
            str(goals_path) if goals_path.exists() else None,
            goal_type=args.goal_type,
        )
        if args.episode_limit is not None:
            dataset.episodes = dataset.episodes[:args.episode_limit]
        attach_resolved_scene_paths(dataset, scene_root)
        print(f"Loaded {len(dataset)} {args.goal_type} episodes for scene {scene_dir.name}.")

        results = evaluate(
            env,
            dataset,
            policy,
            max_steps=args.max_steps,
            threshold=args.threshold,
            log_episodes=args.log_episodes,
            render_live=args.render_live,
            render_delay_ms=args.render_delay_ms,
            instruction_policy=args.instruction_policy,
            seed=args.seed,
            log_prompts=args.log_prompts,
            log_positions=args.log_positions,
            action_protocol=args.action_protocol,
            space_coverage_samples=args.space_coverage_samples,
            space_coverage_radius=args.space_coverage_radius,
        )

        all_episode_results.extend(results["episodes"])

    if not all_episode_results:
        print("No episodes were evaluated.")
        env.close()
        return

    # Aggregate metrics across all scenes
    from metrics import aggregate_metrics
    final_metrics = aggregate_metrics(all_episode_results)

    print("\nFinal aggregated metrics across all scenes:")
    print(json.dumps(final_metrics, indent=2))

    if args.log_episodes:
        print("\nAll episode details:")
        print(json.dumps(all_episode_results, indent=2))

    env.close()


if __name__ == "__main__":
    main()