code/evolvingnav_paper/world.py

"""High-level Habitat action adapter; evaluator-private fields never enter Agent."""

from __future__ import annotations

import habitat_sim
import numpy as np

from evolvingnav_paper.agent import ViewEvidence
from evolvingnav_paper.coverage import (
    camera_forward, camera_transform, candidate_surface_samples, depth_quality,
    heading_quaternion,
    visible_sample_ids,
)
from evolvingnav_paper.habitat_utils import set_agent
from evolvingnav_paper.memory import backproject


class HabitatAgentWorld:
    def __init__(self, backend, viewpoints: dict[int, dict], state_centers: dict[int, list],
                 start_xyz, start_xyzw, *, speed_mps: float = 1.0,
                 calibrator=None, category_recall: float = 0.8,
                 known_states: set[int] | None = None,
                 motion_schedule: list[dict] | None = None,
                 surface_points: dict[int, list] | None = None) -> None:
        self.backend = backend
        self.viewpoints = viewpoints
        self.state_centers = state_centers
        self.position = np.asarray(start_xyz, dtype=float)
        self.rotation = list(start_xyzw)
        self.speed_mps = speed_mps
        self.calibrator = calibrator
        self.category_recall = category_recall
        self.frame = 0
        self.private_inspections: list[dict] = []
        self.last_detection = None
        surface_points = surface_points or {}
        self.samples = {state: candidate_surface_samples(
            center, place_points=surface_points.get(state))
                        for state, center in state_centers.items()}
        self.frontiers = set(viewpoints) - (known_states or set())
        self.motion_schedule = sorted(motion_schedule or [], key=lambda row: row["time_s"])
        self.world_time_s = 0.0
        self._next_motion = 0

    def advance_time(self, seconds: float) -> None:
        if seconds < 0:
            raise ValueError("world time cannot move backwards")
        self.world_time_s += seconds
        while (self._next_motion < len(self.motion_schedule)
               and self.motion_schedule[self._next_motion]["time_s"] <= self.world_time_s):
            self.backend.move_target(self.motion_schedule[self._next_motion])
            self._next_motion += 1

    def distance(self, goal) -> float:
        return self.backend.distance(self.position, goal)

    def sample_count(self, state: int) -> int:
        return len(self.samples[state])

    def expected_new_detection(self, state: int, uncovered: float) -> float:
        center = np.asarray(self.state_centers[state], dtype=float)
        viewpoint = self.viewpoints[state]
        distance = float(np.linalg.norm(center - np.asarray(viewpoint["position_xyz"])))
        direction = center - np.asarray(viewpoint["position_xyz"])
        angle = max(0.0, float(np.dot(
            direction / max(distance, 1e-9),
            camera_forward(viewpoint["rotation_xyzw"]),
        )))
        if self.calibrator is None:
            return self.category_recall * uncovered
        return self.calibrator.predict({
            "coverage": uncovered,
            "range_m": distance,
            "angle_cos": angle,
            "projected_pixels": int(25 * uncovered),
            "depth_quality": 1.0,
            "category_recall": self.category_recall,
        })

    def move_chunk(self, goal, max_distance: float) -> tuple[float, float]:
        path = habitat_sim.ShortestPath()
        path.requested_start = np.asarray(self.position, dtype=np.float32)
        path.requested_end = np.asarray(goal, dtype=np.float32)
        if not self.backend.real.pathfinder.find_path(path):
            return 0.0, 0.0
        remaining = min(float(max_distance), float(path.geodesic_distance))
        points = [np.asarray(point, dtype=float) for point in path.points]
        new_position = points[0]
        for point in points[1:]:
            length = float(np.linalg.norm(point - new_position))
            if length >= remaining:
                new_position = new_position + (point - new_position) * (remaining / length)
                break
            remaining -= length
            new_position = point
        displacement = min(float(max_distance), float(path.geodesic_distance))
        heading = new_position - self.position
        if np.linalg.norm(heading[[0, 2]]) > 1e-6:
            self.rotation = heading_quaternion(heading)
        self.position = new_position
        set_agent(self.backend.real.get_agent(0), self.position.tolist(), self.rotation)
        self.advance_time(displacement / self.speed_mps)
        return displacement, displacement / self.speed_mps

    def inspect(self, state: int) -> tuple[bool, list[ViewEvidence]]:
        viewpoint = self.viewpoints[state]
        self.position = np.asarray(viewpoint["position_xyz"], dtype=float)
        self.rotation = viewpoint["rotation_xyzw"]
        private = self.backend.inspect(state, self.position.tolist())
        self.private_inspections.append({"state_id": state, **private})
        observation = self.backend.last_observation
        if observation is None:
            raise RuntimeError("Habitat did not produce an RGB-D frame")
        self.last_detection = None
        if self.backend.last_detections:
            strongest = max(self.backend.last_detections, key=lambda item: item.confidence)
            depth = observation["depth"]
            pixels = np.argwhere(strongest.mask & np.isfinite(depth) & (depth > 0))
            if len(pixels):
                v, u = np.median(pixels, axis=0)
                depth_m = float(np.median(depth[pixels[:, 0], pixels[:, 1]]))
                height, width = depth.shape
                focal = width / (2 * np.tan(np.deg2rad(79.0) / 2))
                intrinsics = np.array([[focal, 0, width / 2],
                                       [0, focal, height / 2], [0, 0, 1]], dtype=float)
                point = backproject(u, v, depth_m, intrinsics,
                                    camera_transform(self.position, self.rotation))
                self.last_detection = {
                    "world_point": point,
                    "confidence": strongest.confidence,
                    "evidence_id": f"frame-{self.frame + 1}:positive",
                }
        return bool(private["detected"]), self._evidence_from_depth(observation["depth"])

    def observe_chunk(self) -> list[ViewEvidence]:
        if self.backend.detector is None:
            return []
        observation = self.backend.real.get_sensor_observations()
        depth = np.asarray(observation["depth"])
        covered_any = any(visible_sample_ids(
            samples, self.position, self.rotation, depth, 79.0
        ) for samples in self.samples.values())
        if not covered_any:
            return []
        detected = self.backend.detector(
            np.asarray(observation["rgb"]), depth, self.backend.target_category
        )
        return [] if detected else self._evidence_from_depth(depth)

    def _evidence_from_depth(self, depth: np.ndarray) -> list[ViewEvidence]:
        self.frame += 1
        evidence = []
        for candidate, samples in self.samples.items():
            covered = visible_sample_ids(
                samples, self.position, self.rotation, depth, 79.0
            )
            if covered:
                fraction = len(covered) / len(samples)
                if self.calibrator is None:
                    detection_probability = self.category_recall * fraction
                else:
                    direction = np.asarray(self.state_centers[candidate]) - self.position
                    distance = float(np.linalg.norm(direction))
                    features = {
                        "coverage": fraction,
                        "range_m": distance,
                        "angle_cos": max(0.0, float(np.dot(
                            direction / max(distance, 1e-9), camera_forward(self.rotation)))),
                        "projected_pixels": len(covered),
                        "depth_quality": depth_quality(depth),
                        "category_recall": self.category_recall,
                    }
                    detection_probability = self.calibrator.predict(features)
                evidence.append(ViewEvidence(
                    f"frame-{self.frame}:state-{candidate}", candidate, covered,
                    detection_probability, fraction, tuple(self.position.tolist()),
                ))
        return evidence

    def explore(self, budget_m: float) -> tuple[dict[int, tuple[object, float]], float, float]:
        reachable = [
            (self.distance(self.viewpoints[state]["position_xyz"]), state)
            for state in self.frontiers
        ]
        reachable = [(distance, state) for distance, state in reachable
                     if np.isfinite(distance) and distance <= budget_m]
        if not reachable:
            return {}, 0.0, 0.0
        distance, state = min(reachable)
        self.frontiers.remove(state)
        goal = self.viewpoints[state]["position_xyz"]
        displacement, duration = self.move_chunk(goal, distance)
        self.rotation = self.viewpoints[state]["rotation_xyzw"]
        set_agent(self.backend.real.get_agent(0), self.position.tolist(), self.rotation)
        observations = self.backend.real.get_sensor_observations()
        detected = (
            bool(self.backend.detector(
                np.asarray(observations["rgb"]), np.asarray(observations["depth"]),
                self.backend.target_category,
            )) if self.backend.detector is not None else False
        )
        return {state: (goal, 1.0 if detected else 0.5)}, duration, displacement