policy.py

import json
import base64
import io
from pathlib import Path
import numpy as np
from typing import Any, Dict, List, Optional

from PIL import Image
import requests


class RandomPolicy:
    """Simple policy that chooses a random valid action from the environment."""

    def __init__(self, action_space: Any, seed: Optional[int] = None):
        self.action_space = action_space
        self.rng = np.random.default_rng(seed)
        self.last_instruction = None
        self.log_actions = False
        self.log_action_limit = None
        self._logged_actions = 0
        self.save_api_images_dir = None
        self._policy_step_index = 0
        self.current_episode_id = None

    def act(self, rgb: Any = None, depth: Any = None, prompt: Optional[str] = None) -> Any:
        self.last_instruction = prompt
        if hasattr(self.action_space, "n"):
            action = int(self.rng.integers(self.action_space.n))
            self._log_policy_action(action, source="random")
            return action
        if hasattr(self.action_space, "sample"):
            action = self.action_space.sample()
            self._log_policy_action(action, source="random")
            return action
        self._log_policy_action(0, source="random")
        return 0

    def _action_name(self, action: Any) -> str:
        names = {
            0: "finish",
            1: "move_forward",
            2: "turn_left",
            3: "turn_right",
            4: "look_up",
            5: "look_down",
            6: "target_found",
            9: "legacy_finish",
        }
        try:
            return names.get(int(action), str(action))
        except (TypeError, ValueError):
            return str(action)

    def _is_valid_action(self, action: Any) -> bool:
        try:
            action = int(action)
        except (TypeError, ValueError):
            return False
        if action in {6, 9}:
            return True
        if hasattr(self.action_space, "n"):
            return 0 <= action < int(self.action_space.n)
        if hasattr(self.action_space, "contains"):
            return self.action_space.contains(action)
        return True

    def _should_log_action(self) -> bool:
        if not self.log_actions:
            return False
        if self.log_action_limit is None:
            return True
        return self._logged_actions < self.log_action_limit

    def _log_policy_action(
        self,
        action: Any,
        source: str,
        episode: Any = None,
        raw_output: Any = None,
        saved_image_paths: Optional[List[str]] = None,
    ) -> None:
        if not self._should_log_action():
            return
        episode_id = getattr(episode, "episode_id", None) or self.current_episode_id or "unknown"
        message = (
            f"[policy] episode={episode_id} source={source} "
            f"action={action}({self._action_name(action)})"
        )
        if raw_output is not None:
            message += f" raw={raw_output!r}"
        if saved_image_paths:
            message += f" images={saved_image_paths}"
        print(message)
        self._logged_actions += 1


class HumanPolicy(RandomPolicy):
    """Interactive policy controlled from the terminal."""

    def __init__(self, action_space: Any, seed: Optional[int] = None):
        super().__init__(action_space, seed=seed)
        self._last_prompt = None

    def act(self, rgb: Any = None, depth: Any = None, prompt: Optional[str] = None) -> Any:
        if prompt and prompt != self._last_prompt:
            episode_id = self.current_episode_id or "unknown"
            print(f"\n[prompt] episode={episode_id}")
            print(prompt)
            print(
                "Controls: 0=finish episode, 1=move_forward, 2=turn_left, 3=turn_right, "
                "4=look_up, 5=look_down, 6=target_found, 9=legacy_finish, q=quit"
            )
            self._last_prompt = prompt

        while True:
            value = input("Action [0/1/2/3/4/5/6/9/q]: ").strip().lower()
            if value in {"q", "quit", "exit"}:
                raise KeyboardInterrupt("Human control stopped by user.")
            try:
                action = int(value)
            except ValueError:
                print("Please enter 0, 1, 2, 3, 4, 5, 6, 9, or q.")
                continue
            if self._is_valid_action(action):
                self._log_policy_action(action, source="human")
                return action
            print("Action is outside the environment action space.")


class ModelApiPolicy(RandomPolicy):
    """Policy that sends observation and instruction to an external API to select actions."""

    def __init__(
        self,
        action_space: Any,
        api_key: str,
        api_url: str,
        model_name: str,
        seed: Optional[int] = None,
    ):
        super().__init__(action_space, seed=seed)
        self.api_key = api_key
        self.api_url = api_url
        self.model_name = model_name
        self._warned_api_failure = False
        self.last_api_raw_output = None

    def _build_payload(self, rgb: Any = None, depth: Any = None, prompt: Optional[str] = None) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "model": self.model_name,
            "input": {
                "sensors": {},
            },
            "images": [],
        }

        image_item = self._image_content_item(rgb, "rgb")
        if image_item is not None:
            payload["images"].append(image_item)
            payload["input"]["sensors"]["rgb"] = "<included as rgb image>"
        else:
            payload["input"]["sensors"]["rgb"] = "<rgb image unavailable>"

        image_item = self._image_content_item(depth, "depth")
        if image_item is not None:
            payload["images"].append(image_item)
            payload["input"]["sensors"]["depth"] = "<included as depth image>"
        else:
            payload["input"]["sensors"]["depth"] = "<depth image unavailable>"

        payload["prompt"] = self._format_prompt(payload, prompt)
        return payload

    def _episode_instruction(self, episode: Any) -> str:
        natural_language = getattr(episode, "natural_language", None)
        if isinstance(natural_language, str):
            return natural_language
        if natural_language:
            return " ".join(str(item) for item in natural_language)
        return ""

    def _image_content_item(self, value: Any, image_kind: str) -> Optional[Dict[str, Any]]:
        image = self._to_pil_image(value, image_kind)
        if image is None:
            return None

        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
        return {
            "type": "image_url",
            "name": image_kind,
            "image_url": {
                "url": f"data:image/png;base64,{encoded}",
            },
        }

    def _save_api_images(self, payload: Dict[str, Any], episode: Any = None) -> List[str]:
        if self.save_api_images_dir is None:
            return []

        output_dir = Path(self.save_api_images_dir)
        output_dir.mkdir(parents=True, exist_ok=True)
        episode_id = getattr(episode, "episode_id", None) or self.current_episode_id or "unknown"
        step_index = int(self._policy_step_index)

        saved_paths = []
        prefix = "data:image/png;base64,"
        for index, image_item in enumerate(payload.get("images", [])):
            image_url = image_item.get("image_url", {}).get("url", "")
            if not image_url.startswith(prefix):
                continue

            image_kind = image_item.get("name") or f"image{index}"
            filename = f"{episode_id}_step{step_index:04d}_{image_kind}.png"
            path = output_dir / filename
            path.write_bytes(base64.b64decode(image_url[len(prefix):]))
            saved_paths.append(str(path))

        return saved_paths

    def _to_pil_image(self, value: Any, image_kind: str) -> Optional[Image.Image]:
        if not isinstance(value, np.ndarray):
            return None

        array = np.asarray(value)
        if image_kind == "depth":
            return self._depth_to_pil_image(array)
        return self._rgb_to_pil_image(array)

    def _rgb_to_pil_image(self, array: np.ndarray) -> Optional[Image.Image]:
        array = np.squeeze(array)
        if array.ndim != 3:
            return None
        if array.shape[0] in {3, 4} and array.shape[-1] not in {3, 4}:
            array = np.transpose(array, (1, 2, 0))
        if array.shape[-1] == 4:
            array = array[..., :3]
        if array.shape[-1] != 3:
            return None

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

    def _depth_to_pil_image(self, array: np.ndarray) -> Optional[Image.Image]:
        array = np.squeeze(array)
        if array.ndim == 3:
            array = array[..., 0]
        if array.ndim != 2:
            return None

        array = np.nan_to_num(array.astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        valid = array[array > 0.0]
        if valid.size:
            min_depth = float(valid.min())
            max_depth = float(valid.max())
            if max_depth > min_depth:
                array = (array - min_depth) / (max_depth - min_depth)
            elif max_depth <= 1.0:
                array = np.clip(array, 0.0, 1.0)
            else:
                array = np.zeros_like(array)
        array = np.clip(array, 0.0, 1.0)
        array = (array * 255.0).astype(np.uint8)
        return Image.fromarray(array, mode="L")

    def _chat_completions_url(self) -> str:
        api_url = self.api_url.rstrip("/")
        if api_url.endswith("/chat/completions"):
            return api_url
        return f"{api_url}/chat/completions"

    def _format_prompt(self, payload: Dict[str, Any], prompt: Optional[str]) -> str:
        sensors = payload["input"].get("sensors", {})
        action_count = getattr(self.action_space, "n", 4)
        return (
            f"{prompt or 'You are controlling an embodied navigation agent.'}\n\n"
            f"The navigation action space has {action_count} simulator actions plus benchmark signals 6 and 9.\n"
            "Action IDs are: 0=FINISH the whole episode, 1=MOVE_FORWARD by 0.25m, "
            "2=TURN_LEFT by 30 degrees, 3=TURN_RIGHT by 30 degrees, "
            "4=LOOK_UP by 30 degrees, 5=LOOK_DOWN by 30 degrees, "
            "6=TARGET_FOUND when you are near a matching target instance, "
            "9=legacy FINISH alias.\n"
            f"Current non-image observations JSON:\n{json.dumps(sensors, ensure_ascii=False)}\n\n"
            "Return strict JSON only, for example: {\"action\": 1}."
        )

    def _query_api(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
        }
        content: List[Dict[str, Any]] = [{"type": "text", "text": payload["prompt"]}]
        content.extend(payload.get("images", []))

        request_payload = {
            "model": self.model_name,
            "messages": [
                {
                    "role": "system",
                    "content": "You are a navigation policy. Return JSON only.",
                },
                {
                    "role": "user",
                    "content": content,
                },
            ],
            "temperature": 0,
            "response_format": {"type": "json_object"},
        }
        response = requests.post(
            self._chat_completions_url(),
            headers=headers,
            json=request_payload,
            timeout=15,
        )
        response.raise_for_status()
        return response.json()

    def _random_action_without_logging(self) -> Any:
        if hasattr(self.action_space, "n"):
            return int(self.rng.integers(self.action_space.n))
        if hasattr(self.action_space, "sample"):
            return self.action_space.sample()
        return 0

    def _choose_action(self, api_response: Dict[str, Any]) -> Any:
        raw_output = None
        if not isinstance(api_response, dict):
            self.last_api_raw_output = None
            return self._random_action_without_logging()

        output = api_response.get("action")
        raw_output = output
        if output is None and "choices" in api_response:
            try:
                content = api_response["choices"][0]["message"]["content"]
                raw_output = content
                parsed = json.loads(content)
                output = parsed.get("action")
            except (KeyError, IndexError, TypeError, json.JSONDecodeError):
                output = None
        if output is None:
            action_map = api_response.get("actions")
            if isinstance(action_map, dict) and "selected" in action_map:
                output = action_map["selected"]
                raw_output = output

        self.last_api_raw_output = raw_output

        try:
            action = int(output)
            if self._is_valid_action(action):
                return action
        except (TypeError, ValueError):
            pass

        if isinstance(output, str) and output.isdigit():
            action = int(output)
            if self._is_valid_action(action):
                return action

        if isinstance(output, 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(output.strip().lower())
            if action is not None and self._is_valid_action(action):
                return action

        return self._random_action_without_logging()

    def act(self, rgb: Any = None, depth: Any = None, prompt: Optional[str] = None) -> Any:
        self.last_instruction = prompt
        payload = self._build_payload(rgb=rgb, depth=depth, prompt=prompt)
        saved_image_paths = self._save_api_images(payload)
        try:
            api_response = self._query_api(payload)
            action = self._choose_action(api_response)
            self._log_policy_action(
                action,
                source="api",
                raw_output=self.last_api_raw_output,
                saved_image_paths=saved_image_paths,
            )
            self._policy_step_index += 1
            return action
        except Exception as exc:
            if not self._warned_api_failure:
                print(f"Warning: evaluation API request failed; falling back to RandomPolicy. Error: {exc}")
                self._warned_api_failure = True
            action = self._random_action_without_logging()
            self._log_policy_action(
                action,
                source="api_fallback",
                saved_image_paths=saved_image_paths,
            )
            self._policy_step_index += 1
            return action