evaluator.py

import itertools
import numpy as np
from typing import Any, Dict, List, Optional, Sequence

from benchmark_api import Observation
from dataset import CustomDataset, CustomEpisode
from metrics import aggregate_metrics
from prompt import get_navigation_task_prompt


ACTION_ID_TO_NAME = {
    0: "stop",
    1: "move_forward",
    2: "turn_left",
    3: "turn_right",
    4: "look_up",
    5: "look_down",
}

ACTION_PROTOCOLS = {"goat", "legacy", "compat"}
DEFAULT_SPACE_COVERAGE_SAMPLES = 200
DEFAULT_SPACE_COVERAGE_RADIUS = 1.5


def is_near_object(
    agent_position: Any,
    object_id: int,
    target_positions: Dict[int, List[float]],
    threshold: float = 1.0,
) -> bool:
    """Return True when the agent is within threshold meters of the goal position."""
    if int(object_id) not in target_positions:
        return False

    agent_position = np.array(agent_position, dtype=np.float32)
    target_position = np.array(target_positions[int(object_id)], dtype=np.float32)
    distance = np.linalg.norm(agent_position - target_position)
    return distance <= threshold


def required_target_count(mode: str, target_ids: List[int], target_count: Optional[int] = None) -> int:
    mode = str(mode or "any").lower()
    total_targets = len(target_ids or [])
    if mode in {"single", "any"}:
        return 1
    if mode == "many":
        if target_count is None:
            return min(2, total_targets)
        return max(1, min(int(target_count), total_targets))
    if mode == "all":
        return total_targets
    return 1


def check_success(
    mode: str,
    target_ids: List[int],
    reached_ids: List[int],
    target_count: Optional[int] = None,
) -> bool:
    """Return whether the episode succeeded according to the task mode."""
    mode = str(mode or "any").lower()
    target_set = set(int(x) for x in (target_ids or []))
    reached_set = set(int(x) for x in (reached_ids or []))
    if not target_set:
        return False

    if mode in {"single", "any"}:
        return len(target_set & reached_set) > 0
    if mode == "many":
        required = required_target_count(mode, list(target_set), target_count)
        return len(target_set & reached_set) >= required
    if mode == "all":
        return target_set.issubset(reached_set)
    return len(target_set & reached_set) > 0


def _f1_score(precision: float, recall: float) -> float:
    if precision + recall <= 0.0:
        return 0.0
    return 2.0 * precision * recall / (precision + recall)


def _mode_task_type(mode: str) -> str:
    mode = str(mode or "any").lower()
    if mode == "many":
        return "mos"
    if mode == "all":
        return "moc"
    return mode


def _geodesic_distance(pathfinder: Any, start: Sequence[float], end: Sequence[float]) -> Optional[float]:
    try:
        import habitat_sim
    except ModuleNotFoundError:
        return None

    path = habitat_sim.ShortestPath()
    path.requested_start = [float(x) for x in start[:3]]
    path.requested_end = [float(x) for x in end[:3]]
    try:
        found = pathfinder.find_path(path)
    except Exception:
        return None
    if not found or not np.isfinite(path.geodesic_distance):
        return None
    return float(path.geodesic_distance)


def _snap_position(pathfinder: Any, position: Sequence[float]) -> Optional[List[float]]:
    if position is None or len(position) < 3:
        return None
    point = [float(x) for x in position[:3]]
    snap_point = getattr(pathfinder, "snap_point", None)
    if callable(snap_point):
        try:
            snapped = snap_point(point)
            point = [float(snapped[index]) for index in range(3)]
        except Exception:
            return None
    is_navigable = getattr(pathfinder, "is_navigable", None)
    if callable(is_navigable):
        try:
            if not is_navigable(point):
                return None
        except Exception:
            return None
    if not all(np.isfinite(point)):
        return None
    return point


def _episode_nav_positions(
    pathfinder: Any,
    episode: CustomEpisode,
    target_ids: Sequence[int],
) -> Dict[int, List[float]]:
    positions: Dict[int, List[float]] = {}
    for object_id in target_ids:
        int_id = int(object_id)
        snapped = _snap_position(pathfinder, episode.target_positions.get(int_id))
        if snapped is not None:
            positions[int_id] = snapped
    return positions


def _route_distance(
    pathfinder: Any,
    start_position: Sequence[float],
    target_positions: Dict[int, List[float]],
    target_ids: Sequence[int],
) -> Optional[float]:
    requested_ids = [int(object_id) for object_id in target_ids]
    if any(object_id not in target_positions for object_id in requested_ids):
        return None
    ids = list(requested_ids)
    if not ids:
        return 0.0
    if len(ids) != len(set(ids)):
        ids = list(dict.fromkeys(ids))

    start_distances: Dict[int, float] = {}
    for object_id in ids:
        distance = _geodesic_distance(pathfinder, start_position, target_positions[object_id])
        if distance is not None:
            start_distances[object_id] = distance
    if len(start_distances) != len(ids):
        return None

    pair_distances: Dict[tuple[int, int], float] = {}
    for source_id in ids:
        for target_id in ids:
            if source_id == target_id:
                pair_distances[(source_id, target_id)] = 0.0
                continue
            distance = _geodesic_distance(
                pathfinder,
                target_positions[source_id],
                target_positions[target_id],
            )
            if distance is None:
                return None
            pair_distances[(source_id, target_id)] = distance

    # Open-path TSP from the start to all selected targets. Exact DP is cheap for
    # benchmark-sized target sets; greedy fallback prevents pathological blowups.
    if len(ids) > 12:
        remaining = set(ids)
        current_id: Optional[int] = None
        total = 0.0
        while remaining:
            if current_id is None:
                next_id = min(remaining, key=lambda object_id: start_distances[object_id])
                total += start_distances[next_id]
            else:
                next_id = min(remaining, key=lambda object_id: pair_distances[(current_id, object_id)])
                total += pair_distances[(current_id, next_id)]
            remaining.remove(next_id)
            current_id = next_id
        return total

    dp: Dict[tuple[int, int], float] = {}
    for index, object_id in enumerate(ids):
        dp[(1 << index, index)] = start_distances[object_id]

    full_mask = (1 << len(ids)) - 1
    for _ in range(1, len(ids)):
        next_dp: Dict[tuple[int, int], float] = {}
        for (mask, last_index), distance in dp.items():
            last_id = ids[last_index]
            for next_index, next_id in enumerate(ids):
                if mask & (1 << next_index):
                    continue
                next_mask = mask | (1 << next_index)
                next_distance = distance + pair_distances[(last_id, next_id)]
                state = (next_mask, next_index)
                if state not in next_dp or next_distance < next_dp[state]:
                    next_dp[state] = next_distance
        dp = next_dp
    candidates = [distance for (mask, _), distance in dp.items() if mask == full_mask]
    return min(candidates) if candidates else None


def _optimal_route_distance(
    pathfinder: Any,
    start_position: Sequence[float],
    target_positions: Dict[int, List[float]],
    target_ids: Sequence[int],
    target_count: int,
) -> Optional[float]:
    ids = [int(object_id) for object_id in target_ids if int(object_id) in target_positions]
    if target_count <= 0:
        return 0.0
    if len(ids) < target_count:
        return None
    if len(ids) <= 12:
        best_distance = None
        for subset in itertools.combinations(ids, target_count):
            distance = _route_distance(pathfinder, start_position, target_positions, subset)
            if distance is None:
                continue
            if best_distance is None or distance < best_distance:
                best_distance = distance
        return best_distance

    remaining = set(ids)
    current_position = [float(x) for x in start_position[:3]]
    total = 0.0
    while remaining and len(ids) - len(remaining) < target_count:
        best_id = None
        best_distance = None
        for object_id in remaining:
            distance = _geodesic_distance(pathfinder, current_position, target_positions[object_id])
            if distance is None:
                continue
            if best_distance is None or distance < best_distance:
                best_distance = distance
                best_id = object_id
        if best_id is None or best_distance is None:
            return None
        total += best_distance
        current_position = target_positions[best_id]
        remaining.remove(best_id)
    return total


def _sample_navigable_points(
    pathfinder: Any,
    sample_count: int,
) -> List[List[float]]:
    if pathfinder is None or sample_count <= 0:
        return []
    get_point = getattr(pathfinder, "get_random_navigable_point", None)
    if not callable(get_point):
        return []

    points: List[List[float]] = []
    seen = set()
    max_attempts = max(sample_count * 10, sample_count)
    for _ in range(max_attempts):
        try:
            raw_point = get_point()
            point = [float(raw_point[index]) for index in range(3)]
        except Exception:
            continue
        if not all(np.isfinite(point)):
            continue
        key = tuple(round(value, 2) for value in point)
        if key in seen:
            continue
        seen.add(key)
        points.append(point)
        if len(points) >= sample_count:
            break
    return points


def _space_exhausted_step(
    trajectory_positions: Sequence[tuple[int, Sequence[float]]],
    coverage_points: Sequence[Sequence[float]],
    coverage_radius: float,
) -> Optional[int]:
    if not coverage_points:
        return None
    remaining = {
        index
        for index, _ in enumerate(coverage_points)
    }
    radius = max(float(coverage_radius), 0.0)
    for step, position in trajectory_positions:
        current = np.asarray(position, dtype=np.float32)
        covered = []
        for index in remaining:
            point = np.asarray(coverage_points[index], dtype=np.float32)
            if float(np.linalg.norm(current - point)) <= radius:
                covered.append(index)
        for index in covered:
            remaining.remove(index)
        if not remaining:
            return int(step)
    return None


def _near_target_ids(
    agent_position: Any,
    target_object_ids: List[int],
    target_positions: Dict[int, List[float]],
    threshold: float,
) -> List[int]:
    return [
        object_id
        for object_id in target_object_ids
        if is_near_object(
            agent_position,
            object_id,
            target_positions,
            threshold=threshold,
        )
    ]


def _reset_environment(env: Any, episode: CustomEpisode, dataset: CustomDataset) -> Any:
    try:
        observations = env.reset(episode=episode)
    except TypeError:
        if hasattr(env, "_dataset"):
            env._dataset = dataset
        if hasattr(env, "_current_episode"):
            env._current_episode = episode
        elif hasattr(env, "current_episode"):
            env.current_episode = episode
        observations = env.reset()
    _ensure_agent_sensor_alias(env)
    return observations


def _ensure_agent_sensor_alias(env: Any) -> None:
    sim = getattr(env, "sim", None)
    agents = getattr(sim, "agents", None)
    if not agents:
        return
    for agent in agents:
        if not hasattr(agent, "sensors") and hasattr(agent, "_sensors"):
            try:
                setattr(agent, "sensors", getattr(agent, "_sensors"))
            except AttributeError:
                pass


def _sensor(observations: Any, name: str) -> Any:
    if isinstance(observations, dict):
        return observations.get(name)
    return None


def _reset_policy(policy: Any) -> None:
    reset = getattr(policy, "reset", None)
    if callable(reset):
        reset()


def _set_policy_episode_context(policy: Any, episode: CustomEpisode) -> None:
    try:
        setattr(policy, "current_episode_id", episode.episode_id)
    except AttributeError:
        pass


def _select_natural_language(
    natural_language: Any,
    instruction_policy: str,
    rng: Optional[np.random.Generator],
) -> List[str]:
    if isinstance(natural_language, str):
        candidates = [natural_language]
    else:
        candidates = list(natural_language or [])

    candidates = [str(item) for item in candidates if str(item).strip()]
    if not candidates:
        return []

    instruction_policy = str(instruction_policy or "first").lower()
    if instruction_policy == "all":
        return candidates
    if instruction_policy == "random":
        if rng is None:
            rng = np.random.default_rng()
        return [candidates[int(rng.integers(len(candidates)))]]
    if instruction_policy == "first":
        return [candidates[0]]
    raise ValueError(f"Unknown instruction policy: {instruction_policy}")


def _episode_prompt(
    episode: CustomEpisode,
    max_steps: int,
    instruction_policy: str,
    rng: Optional[np.random.Generator],
) -> str:
    return get_navigation_task_prompt(
        target_object_ids=episode.target_object_ids,
        target_mode=episode.target_mode,
        target_count=episode.target_count,
        natural_language=_select_natural_language(
            episode.natural_language,
            instruction_policy=instruction_policy,
            rng=rng,
        ),
        goal_type=episode.goal_type,
        goal_text=_goal_text_for_observation(episode, ""),
        max_actions=max_steps,
    )


def _quaternion_xyzw(rotation: Any) -> Optional[tuple[float, float, float, float]]:
    if rotation is None:
        return None

    if hasattr(rotation, "imag") and hasattr(rotation, "real"):
        imag = np.asarray(rotation.imag, dtype=np.float64).reshape(-1)
        if imag.size >= 3:
            return float(imag[0]), float(imag[1]), float(imag[2]), float(rotation.real)

    vector = getattr(rotation, "vector", None)
    scalar = getattr(rotation, "scalar", None)
    if vector is not None and scalar is not None:
        vector = np.asarray(vector, dtype=np.float64).reshape(-1)
        if vector.size >= 3:
            return float(vector[0]), float(vector[1]), float(vector[2]), float(scalar)

    try:
        values = np.asarray(rotation, dtype=np.float64).reshape(-1)
    except (TypeError, ValueError):
        return None
    if values.size != 4:
        return None
    return float(values[0]), float(values[1]), float(values[2]), float(values[3])


def _agent_compass(rotation: Any) -> float:
    """Return an approximate yaw angle in radians from an agent rotation."""
    components = _quaternion_xyzw(rotation)
    if components is None:
        return 0.0
    x, y, z, w = components
    return float(np.arctan2(2.0 * (w * y + x * z), 1.0 - 2.0 * (y * y + z * z)))


def _build_agent_observation(
    observations: Any,
    env: Any,
    episode: CustomEpisode,
    prompt: str,
    step_count: int,
    max_steps: int,
    previous_action: Optional[int],
) -> Observation:
    agent_state = env.sim.get_agent_state()
    position = np.asarray(agent_state.position, dtype=np.float32)
    if position.size >= 3:
        gps = np.asarray([position[0], position[2]], dtype=np.float32)
    else:
        gps = np.zeros(2, dtype=np.float32)

    rgb = _sensor(observations, "rgb")
    if rgb is None:
        rgb = np.zeros((0, 0, 3), dtype=np.uint8)

    goal_text = _goal_text_for_observation(episode, prompt)
    return Observation(
        rgb=np.asarray(rgb),
        depth=_sensor(observations, "depth"),
        goal_text=goal_text,
        goal_embedding=None,
        goal_image=None,
        goal_type=_goal_type_for_episode(episode),
        subtask_type=_goal_type_for_episode(episode),
        target_mode=str(episode.target_mode or "any"),
        target_count=episode.target_count,
        episode_id=episode.episode_id,
        scene_id=episode.scene_id,
        step_count=step_count,
        gps=gps,
        compass=_agent_compass(getattr(agent_state, "rotation", None)),
        max_steps=max_steps,
        task_prompt=prompt,
        previous_action=previous_action,
        extras={
            "task_prompt": prompt,
            "query_program": episode.query_program,
            "goal_embedding": None,
            "goal_image": None,
            "goal_type": _goal_type_for_episode(episode),
            "subtask_type": _goal_type_for_episode(episode),
            "goat_action_protocol": "finish=0,target_found=6,legacy_finish=9",
        },
    )


def _goal_text_for_observation(episode: CustomEpisode, prompt: str) -> str:
    if str(getattr(episode, "goal_type", "") or "").lower() == "object":
        goal_text = str(getattr(episode, "goal_text", "") or "").strip()
        if goal_text:
            return goal_text
        object_category = getattr(episode, "object_category", None)
        if object_category:
            return str(object_category).strip()

    natural_language = episode.natural_language
    if isinstance(natural_language, str) and natural_language.strip():
        return natural_language.strip()
    if natural_language:
        values = [str(item).strip() for item in natural_language if str(item).strip()]
        if values:
            return values[0]
    query_program = episode.query_program or {}
    conditions = query_program.get("where") or []
    labels = []
    for condition in conditions:
        value = condition.get("value")
        if value is not None:
            labels.append(str(value))
    if labels:
        return " ".join(labels)
    return prompt


def _goal_type_for_episode(episode: CustomEpisode) -> str:
    goal_type = str(getattr(episode, "goal_type", "") or "").lower()
    if goal_type in {"object", "description", "image"}:
        return goal_type
    return "description"


def _policy_uses_observation_api(policy: Any) -> bool:
    return bool(getattr(policy, "_benchmark_observation_api", False))


def _policy_act(
    policy: Any,
    agent_observation: Observation,
    rgb: Any,
    depth: Any,
    prompt: str,
) -> Any:
    if _policy_uses_observation_api(policy):
        return policy.act(agent_observation)
    return policy.act(rgb=rgb, depth=depth, prompt=prompt)


def _prompt_with_progress(prompt: str, episode: CustomEpisode, reached_ids: set) -> str:
    mode = str(episode.target_mode or "any").lower()
    if mode not in {"many", "all"}:
        return prompt
    return (
        f"{prompt}\n"
        "Multi-target completion rule:\n"
        "- Output 6 when you believe you are next to a matching target object.\n"
        "- Output 0 only when you believe the task is complete and the episode should end.\n"
        "- Output 9 is accepted as a legacy finish alias.\n"
        "- You will not be told how many targets remain; decide from your own observations.\n"
    )


def _coerce_action(action: Any, action_space: Any) -> Any:
    if isinstance(action, dict):
        action = action.get("action")

    if isinstance(action, str):
        action_names = {
            "stop": 0,
            "move_forward": 1,
            "forward": 1,
            "turn_left": 2,
            "left": 2,
            "turn_right": 3,
            "right": 3,
            "look_up": 4,
            "up": 4,
            "look_down": 5,
            "down": 5,
            "target_found": 6,
            "found": 6,
            "subtask_stop": 6,
            "confirm": 6,
            "finish": 0,
            "done": 0,
            "complete": 0,
            "end": 0,
            "legacy_finish": 9,
        }
        action = action_names.get(action.strip().lower(), action)

    try:
        action = int(action)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"Agent returned an invalid action: {action!r}") from exc

    if action in {6, 9}:
        return action

    if hasattr(action_space, "n"):
        if not 0 <= action < int(action_space.n):
            raise ValueError(f"Agent returned action {action}, outside action space size {action_space.n}.")
    elif hasattr(action_space, "contains") and not action_space.contains(action):
        raise ValueError(f"Agent returned action {action}, which is not accepted by the action space.")
    return action


def _action_role(action: int, mode: str, action_protocol: str) -> str:
    mode = str(mode or "any").lower()
    action_protocol = str(action_protocol or "goat").lower()
    if action_protocol not in ACTION_PROTOCOLS:
        raise ValueError(f"Unknown action protocol: {action_protocol}")

    if action == 6:
        return "target_found"
    if action == 9:
        return "finish"
    if action == 0:
        if mode in {"single", "any"}:
            return "target_found"
        if action_protocol == "goat":
            return "finish"
        return "target_found"
    return "navigate"


def _env_action(action: int) -> str:
    if action not in ACTION_ID_TO_NAME:
        raise ValueError(f"Unknown action id: {action}")
    return ACTION_ID_TO_NAME[action]


def _rgb_frame(rgb: Any) -> Optional[np.ndarray]:
    if rgb is None:
        return None

    frame = np.asarray(rgb)
    frame = np.squeeze(frame)
    if frame.ndim != 3:
        return None
    if frame.shape[0] in {3, 4} and frame.shape[-1] not in {3, 4}:
        frame = np.transpose(frame, (1, 2, 0))
    if frame.shape[-1] == 4:
        frame = frame[..., :3]
    if frame.shape[-1] != 3:
        return None

    if np.issubdtype(frame.dtype, np.floating):
        max_value = float(np.nanmax(frame)) if frame.size else 1.0
        if max_value <= 1.0:
            frame = frame * 255.0
    frame = np.nan_to_num(frame, nan=0.0, posinf=255.0, neginf=0.0)
    return np.clip(frame, 0, 255).astype(np.uint8)


def _show_live_frame(rgb: Any, window_name: str, delay_ms: int) -> bool:
    try:
        import cv2
    except ModuleNotFoundError as exc:
        raise RuntimeError("Live rendering requires OpenCV. Install it with `pip install opencv-python`.") from exc

    frame = _rgb_frame(rgb)
    if frame is None:
        return True

    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    cv2.imshow(window_name, frame)
    key = cv2.waitKey(max(1, delay_ms)) & 0xFF
    return key not in {ord("q"), 27}


def _format_position(position: Any) -> str:
    values = np.asarray(position, dtype=np.float32).tolist()
    return "[" + ", ".join(f"{value:.3f}" for value in values) + "]"


def evaluate(
    env: Any,
    dataset: CustomDataset,
    policy: Any,
    max_steps: int = 500,
    threshold: float = 1.0,
    log_episodes: bool = False,
    render_live: bool = False,
    render_delay_ms: int = 30,
    instruction_policy: str = "first",
    seed: Optional[int] = None,
    log_prompts: bool = False,
    log_positions: bool = False,
    action_protocol: str = "goat",
    space_coverage_samples: int = DEFAULT_SPACE_COVERAGE_SAMPLES,
    space_coverage_radius: float = DEFAULT_SPACE_COVERAGE_RADIUS,
) -> Dict[str, Any]:
    """Evaluate a policy on a custom multi-object navigation dataset."""
    action_protocol = str(action_protocol or "goat").lower()
    if action_protocol not in ACTION_PROTOCOLS:
        raise ValueError(f"Unknown action protocol: {action_protocol}")
    episode_results: List[Dict[str, Any]] = []
    instruction_rng = np.random.default_rng(seed)

    for episode_index, episode in enumerate(dataset.episodes):
        observations = _reset_environment(env, episode, dataset)
        _set_policy_episode_context(policy, episode)
        _reset_policy(policy)
        prompt = _episode_prompt(
            episode,
            max_steps,
            instruction_policy=instruction_policy,
            rng=instruction_rng,
        )
        if log_prompts:
            print(f"\n[prompt] episode={episode.episode_id}")
            print(prompt)
        agent_state = env.sim.get_agent_state()
        previous_position = np.array(agent_state.position, dtype=np.float32)
        episode_start_position = previous_position.copy()
        trajectory_positions: List[tuple[int, List[float]]] = [
            (0, previous_position.astype(float).tolist())
        ]
        pathfinder = getattr(getattr(env, "sim", None), "pathfinder", None)
        space_coverage_points = _sample_navigable_points(
            pathfinder,
            int(space_coverage_samples),
        )
        if log_positions:
            print(
                f"[position] episode={episode.episode_id} "
                f"start={_format_position(previous_position)}"
            )

        reached_ids = set()
        trajectory_length = 0.0
        success = False
        failed = False
        stop_count = 0
        finish_count = 0
        true_positive_count = 0
        false_positive_count = 0
        episode_finished = False
        failure_reason = None
        previous_action: Optional[int] = None
        report_events: List[Dict[str, Any]] = []
        target_report_sequence: List[int] = []
        unique_report_sequence: List[int] = []
        finish_step: Optional[int] = None
        all_required_found_step: Optional[int] = None

        for step in range(max_steps):
            if getattr(env, "episode_over", False):
                failure_reason = failure_reason or "environment_episode_over"
                episode_finished = True
                break

            if render_live:
                keep_running = _show_live_frame(
                    _sensor(observations, "rgb"),
                    "Habitat Agent First-Person View",
                    render_delay_ms,
                )
                if not keep_running:
                    env.close()
                    raise KeyboardInterrupt("Live rendering stopped by user.")

            prompt_for_step = _prompt_with_progress(prompt, episode, reached_ids)
            agent_observation = _build_agent_observation(
                observations,
                env,
                episode,
                prompt_for_step,
                step_count=step,
                max_steps=max_steps,
                previous_action=previous_action,
            )
            action = _policy_act(
                policy,
                agent_observation,
                rgb=_sensor(observations, "rgb"),
                depth=_sensor(observations, "depth"),
                prompt=prompt_for_step,
            )
            action = _coerce_action(action, env.action_space)
            previous_action = int(action)
            mode = str(episode.target_mode or "any").lower()
            action_role = _action_role(int(action), mode, action_protocol)

            agent_state = env.sim.get_agent_state()
            previous_step_position = np.array(agent_state.position, dtype=np.float32)
            trajectory_positions.append((step, previous_step_position.astype(float).tolist()))

            if action_role == "target_found":
                current_position = previous_step_position
                stop_count += 1
                near_targets = _near_target_ids(
                    current_position,
                    episode.target_object_ids,
                    episode.target_positions,
                    threshold=threshold,
                )

                mode = str(episode.target_mode or "any").lower()
                if mode in {"single", "any"}:
                    new_targets = [object_id for object_id in near_targets if object_id not in reached_ids]
                    reached_ids.update(near_targets)
                    if near_targets:
                        true_positive_count = 1
                        target_report_sequence.extend(sorted(int(object_id) for object_id in near_targets))
                        unique_report_sequence.extend(sorted(int(object_id) for object_id in new_targets))
                    else:
                        false_positive_count = 1
                        failure_reason = "wrong_stop"
                    report_events.append(
                        {
                            "step": step,
                            "position": current_position.astype(float).tolist(),
                            "near_targets": sorted(int(object_id) for object_id in near_targets),
                            "new_targets": sorted(int(object_id) for object_id in new_targets),
                            "is_true_positive": bool(near_targets),
                            "is_false_positive": not bool(near_targets),
                        }
                    )
                    if log_positions:
                        print(
                            f"[stop] episode={episode.episode_id} "
                            f"near={sorted(near_targets)}"
                        )
                    episode_finished = True
                    break

                if mode in {"many", "all"}:
                    new_targets = [object_id for object_id in near_targets if object_id not in reached_ids]
                    duplicate_targets = [
                        object_id for object_id in near_targets if object_id in reached_ids
                    ]
                    target_report_sequence.extend(sorted(int(object_id) for object_id in near_targets))
                    report_events.append(
                        {
                            "step": step,
                            "position": current_position.astype(float).tolist(),
                            "near_targets": sorted(int(object_id) for object_id in near_targets),
                            "new_targets": sorted(int(object_id) for object_id in new_targets),
                            "duplicate_targets": sorted(int(object_id) for object_id in duplicate_targets),
                            "is_true_positive": bool(new_targets),
                            "is_false_positive": not bool(new_targets),
                        }
                    )
                    if not new_targets:
                        false_positive_count += 1
                        if log_positions:
                            print(
                                f"[confirm] episode={episode.episode_id} "
                                f"near={sorted(near_targets)} new=[] false_positive=True"
                            )
                        continue

                    reached_ids.update(new_targets)
                    unique_report_sequence.extend(sorted(int(object_id) for object_id in new_targets))
                    true_positive_count += len(new_targets)
                    required_count_for_step = required_target_count(
                        episode.target_mode,
                        episode.target_object_ids,
                        target_count=episode.target_count,
                    )
                    if (
                        all_required_found_step is None
                        and len(reached_ids) >= required_count_for_step
                    ):
                        all_required_found_step = step
                    if log_positions:
                        print(
                            f"[confirm] episode={episode.episode_id} "
                            f"new={sorted(new_targets)} reached={sorted(reached_ids)} "
                        )
                    continue

                failed = True
                failure_reason = "unsupported_target_mode"
                episode_finished = True
                break

            if action_role == "finish":
                finish_count += 1
                finish_step = step
                mode = str(episode.target_mode or "any").lower()
                if mode in {"many", "all"}:
                    if log_positions:
                        print(
                            f"[finish] episode={episode.episode_id} "
                            f"reached={sorted(reached_ids)}"
                        )
                    episode_finished = True
                    break

                failed = True
                failure_reason = "finish_not_supported_for_single_or_any"
                episode_finished = True
                break

            observations = env.step(_env_action(action))

            agent_state = env.sim.get_agent_state()
            current_position = np.array(agent_state.position, dtype=np.float32)
            trajectory_length += float(np.linalg.norm(current_position - previous_position))
            trajectory_positions.append((step + 1, current_position.astype(float).tolist()))
            if log_positions:
                print(
                    f"[position] episode={episode.episode_id} step={step + 1} "
                    f"action={action} position={_format_position(current_position)}"
                )
            previous_position = current_position

            if getattr(env, "episode_over", False):
                failure_reason = failure_reason or "environment_episode_over"
                episode_finished = True
                break

        if not episode_finished and not failed and failure_reason is None:
            failure_reason = "max_steps"
        if failure_reason is not None:
            failed = True

        required_count = required_target_count(
            episode.target_mode,
            episode.target_object_ids,
            target_count=episode.target_count,
        )
        mode = str(episode.target_mode or "any").lower()
        task_type = _mode_task_type(mode)
        recall = 0.0
        if required_count:
            recall = min(1.0, len(reached_ids) / required_count)

        predicted_positive_count = true_positive_count + false_positive_count
        precision = 0.0
        if predicted_positive_count:
            precision = true_positive_count / predicted_positive_count
        f1 = _f1_score(precision, recall)

        finished_by_agent = finish_count > 0
        if mode in {"single", "any"}:
            success = len(reached_ids) > 0 and false_positive_count == 0
        elif mode == "many":
            success = len(reached_ids) >= required_count
        elif mode == "all":
            success = len(reached_ids) >= required_count and finished_by_agent
        else:
            success = False
        if success:
            failed = False
            if failure_reason == "max_steps" and mode == "many":
                failure_reason = None
        coverage = 0.0
        if required_count:
            coverage = recall

        found_count = len(reached_ids)
        total_targets = len(episode.target_object_ids)
        found_fraction = 0.0
        if total_targets:
            found_fraction = min(1.0, found_count / total_targets)

        target_positions = (
            _episode_nav_positions(pathfinder, episode, episode.target_object_ids)
            if pathfinder is not None
            else {}
        )
        selected_target_ids = sorted(int(object_id) for object_id in reached_ids)
        if mode == "many" and required_count > 0:
            selected_target_ids = list(dict.fromkeys(unique_report_sequence))[:required_count]

        agent_subset_distance = None
        if pathfinder is not None and selected_target_ids:
            agent_subset_distance = _route_distance(
                pathfinder,
                episode_start_position,
                target_positions,
                selected_target_ids,
            )

        optimal_subset_distance = float(episode.shortest_path_distance or 0.0)
        if optimal_subset_distance <= 0.0 and pathfinder is not None:
            computed_optimal = _optimal_route_distance(
                pathfinder,
                episode_start_position,
                target_positions,
                episode.target_object_ids,
                required_count,
            )
            optimal_subset_distance = float(computed_optimal or 0.0)

        raw_spl_multi = 0.0
        if optimal_subset_distance > 0.0:
            raw_spl_multi = optimal_subset_distance / max(
                optimal_subset_distance,
                trajectory_length,
                1e-6,
            )
        spl_multi = float(success) * raw_spl_multi

        pe = 0.0
        if agent_subset_distance is not None and agent_subset_distance > 0.0:
            pe = agent_subset_distance / max(agent_subset_distance, trajectory_length, 1e-6)

        if mode == "many" and agent_subset_distance is not None and agent_subset_distance > 0.0:
            tse = min(1.0, optimal_subset_distance / agent_subset_distance)
        else:
            tse = 1.0 if agent_subset_distance is not None or mode in {"single", "any", "all"} else 0.0

        target_report_count = len(target_report_sequence)
        rtsr = 0.0
        if target_report_count:
            rtsr = (target_report_count - len(set(target_report_sequence))) / target_report_count

        space_exhausted_step = _space_exhausted_step(
            trajectory_positions,
            space_coverage_points,
            space_coverage_radius,
        )
        space_oracle_enabled = bool(space_coverage_points)
        space_exhausted = space_exhausted_step is not None
        false_termination = bool(
            mode == "all"
            and finished_by_agent
            and (
                len(reached_ids) < required_count
                or (space_oracle_enabled and not space_exhausted)
            )
        )
        termination_diagnosis = "false_termination" if false_termination else None
        if false_termination and not success and failure_reason is None:
            failure_reason = "false_termination"
            failed = True
        elif not success and finished_by_agent and failure_reason is None:
            failure_reason = "incomplete_finish"
            failed = True
        over_exploration_target = 0.0
        if all_required_found_step is not None:
            stop_step = finish_step if finish_step is not None else max_steps
            over_exploration_target = max(0.0, stop_step - all_required_found_step) / max(1, max_steps)
        over_exploration_space = None
        if space_exhausted_step is not None:
            stop_step = finish_step if finish_step is not None else max_steps
            over_exploration_space = max(0.0, stop_step - space_exhausted_step) / max(1, max_steps)

        episode_result = {
            "episode_id": episode.episode_id,
            "goal_type": _goal_type_for_episode(episode),
            "goal_text": _goal_text_for_observation(episode, ""),
            "object_category": getattr(episode, "object_category", None),
            "target_mode": episode.target_mode,
            "task_type": task_type,
            "success": success,
            "sr": float(success),
            "precision": precision,
            "recall": recall,
            "f1": f1,
            "coverage": coverage,
            "spl_multi": spl_multi,
            "spl": spl_multi,
            "raw_spl_multi": raw_spl_multi,
            "raw_spl": raw_spl_multi,
            "pe": pe,
            "tse": tse,
            "rtsr": rtsr,
            "false_termination": float(false_termination),
            "termination_diagnosis": termination_diagnosis,
            "oer_target": over_exploration_target,
            "oer_space": over_exploration_space,
            "space_oracle_enabled": space_oracle_enabled,
            "space_exhausted": space_exhausted,
            "space_exhausted_step": space_exhausted_step,
            "space_coverage_sample_count": len(space_coverage_points),
            "space_coverage_radius": float(space_coverage_radius),
            "trajectory_length": trajectory_length,
            "shortest_path_distance": optimal_subset_distance,
            "optimal_subset_distance": optimal_subset_distance,
            "agent_subset_distance": agent_subset_distance,
            "selected_targets_for_efficiency": selected_target_ids,
            "reached_targets": sorted(list(reached_ids)),
            "found_count": found_count,
            "total_targets": total_targets,
            "found_fraction": found_fraction,
            "target_count": episode.target_count,
            "required_target_count": required_count,
            "stop_count": stop_count,
            "finish_count": finish_count,
            "finished_by_agent": finished_by_agent,
            "action_protocol": action_protocol,
            "finish_step": finish_step,
            "all_required_found_step": all_required_found_step,
            "true_positive_count": true_positive_count,
            "false_positive_count": false_positive_count,
            "report_events": report_events,
            "target_report_sequence": target_report_sequence,
            "failed": failed,
            "failure_reason": failure_reason,
            "natural_language": episode.natural_language,
            "query_program": episode.query_program,
        }
        episode_results.append(episode_result)

        if log_episodes:
            print(
                f"Episode {episode_index + 1}/{len(dataset)} "
                f"id={episode.episode_id} mode={episode.target_mode} "
                f"success={success} "
                f"precision={precision:.3f} recall={recall:.3f} f1={f1:.3f} "
                f"spl_multi={spl_multi:.3f} "
                f"reached={episode_result['reached_targets']} "
                f"failure_reason={failure_reason}"
            )

    metrics = aggregate_metrics(episode_results)
    if render_live:
        try:
            import cv2

            cv2.destroyAllWindows()
        except ModuleNotFoundError:
            pass
    return {
        "metrics": metrics,
        "episodes": episode_results,
    }