benchmark_api.py

from __future__ import annotations

from dataclasses import dataclass
from enum import IntEnum
from typing import Any, Optional, Protocol, runtime_checkable

import numpy as np


class Action(IntEnum):
    """Benchmark action ids returned by agents."""

    FINISH = 0
    STOP = 0
    MOVE_FORWARD = 1
    TURN_LEFT = 2
    TURN_RIGHT = 3
    LOOK_UP = 4
    LOOK_DOWN = 5
    TARGET_FOUND = 6
    SUBTASK_STOP = 6
    LEGACY_FINISH = 9


@dataclass(frozen=True)
class Observation:
    """Observation passed to custom benchmark agents.

    Custom agents loaded with `run_eval.py --agent module:ClassName` receive one
    instance of this dataclass at every step and should return an int action id.
    """

    rgb: np.ndarray
    depth: Optional[np.ndarray]
    goal_text: str
    goal_embedding: Optional[np.ndarray]
    goal_image: Optional[np.ndarray]
    goal_type: str
    subtask_type: str
    target_mode: str
    target_count: Optional[int]
    episode_id: str
    scene_id: str
    step_count: int
    gps: np.ndarray
    compass: float
    max_steps: int
    task_prompt: Optional[str] = None
    previous_action: Optional[int] = None
    extras: Optional[dict[str, Any]] = None


@runtime_checkable
class AgentProtocol(Protocol):
    """Minimal interface for a custom navigation agent."""

    def act(self, observation: Observation) -> int:
        """Return one benchmark action id."""
        ...

    def reset(self) -> None:
        """Optional episode-level state reset."""
        ...