synth-ai 0.2.4.dev4__py3-none-any.whl → 0.2.4.dev6__py3-none-any.whl

synth_ai/environments/examples/nethack/helpers/observation_utils.py

@@ -0,0 +1,433 @@
+"""Observation processing utilities for NetHack."""
+
+from typing import Dict, Any, List, Optional, Tuple
+import re
+
+
+def format_observation_for_llm(observation: Dict[str, Any]) -> str:
+    """
+    Format NetHack observation for LLM consumption.
+
+    Args:
+        observation: Raw observation dictionary
+
+    Returns:
+        Formatted string suitable for LLM input
+    """
+    lines = []
+
+    # Header with turn count and location
+    lines.append(f"=== NetHack - Turn {observation.get('turn_count', 0)} ===")
+    lines.append(f"Dungeon Level: {observation.get('dungeon_level', 1)}")
+
+    # ASCII map
+    if "ascii_map" in observation:
+        lines.append("\n--- Dungeon Map ---")
+        lines.append(observation["ascii_map"])
+
+    # Game message
+    if "message" in observation and observation["message"]:
+        # Clean null bytes and trailing whitespace
+        message = observation["message"].rstrip("\x00").strip()
+        if message:
+            lines.append(f"\nMessage: {message}")
+
+    # Character stats
+    if "character_stats" in observation:
+        stats = observation["character_stats"]
+        lines.append("\n--- Character Status ---")
+        lines.append(f"HP: {stats.get('hp', 0)}/{stats.get('max_hp', 0)}")
+        lines.append(f"Level: {stats.get('level', 1)} (Exp: {stats.get('experience', 0)})")
+        lines.append(f"AC: {stats.get('ac', 10)}, Gold: {stats.get('gold', 0)}")
+
+        # Attributes
+        attrs = []
+        for attr in [
+            "strength",
+            "dexterity",
+            "constitution",
+            "intelligence",
+            "wisdom",
+            "charisma",
+        ]:
+            if attr in stats:
+                attrs.append(f"{attr[:3].upper()}:{stats[attr]}")
+        if attrs:
+            lines.append(f"Attributes: {' '.join(attrs)}")
+
+    # Inventory summary
+    if "inventory_summary" in observation:
+        lines.append("\n--- Inventory ---")
+        lines.append(observation["inventory_summary"])
+
+    # Menu items if in menu
+    if observation.get("in_menu", False) and "menu_items" in observation:
+        lines.append("\n--- Menu Options ---")
+        for i, item in enumerate(observation["menu_items"]):
+            if i < 26:
+                lines.append(f"{chr(ord('a') + i)}) {item}")
+            else:
+                lines.append(f"{i - 26}) {item}")
+
+    # Score and rewards
+    lines.append("\n--- Progress ---")
+    lines.append(f"Score: {observation.get('score', 0)}")
+    lines.append(f"Total Reward: {observation.get('total_reward', 0.0):.2f}")
+    lines.append(f"Last Reward: {observation.get('reward_last', 0.0):.2f}")
+
+    # Termination status
+    if observation.get("terminated", False):
+        lines.append("\n*** GAME OVER ***")
+
+    return "\n".join(lines)
+
+
+def parse_ascii_map(ascii_map: str) -> Dict[str, Any]:
+    """
+    Parse ASCII map to extract key information.
+
+    Args:
+        ascii_map: Raw ASCII map string
+
+    Returns:
+        Dictionary with extracted map information
+    """
+    lines = ascii_map.strip().split("\n")
+
+    map_info = {
+        "width": 0,
+        "height": len(lines),
+        "player_position": None,
+        "stairs_positions": [],
+        "door_positions": [],
+        "item_positions": [],
+        "monster_positions": [],
+        "wall_positions": [],
+        "floor_positions": [],
+    }
+
+    if lines:
+        map_info["width"] = max(len(line) for line in lines)
+
+    # Common NetHack ASCII symbols
+    symbols = {
+        "@": "player",
+        "<": "stairs_up",
+        ">": "stairs_down",
+        "+": "closed_door",
+        "-": "open_door_horizontal",
+        "|": "open_door_vertical",
+        ".": "floor",
+        "#": "corridor",
+        " ": "wall",
+        "$": "gold",
+        "*": "gem",
+        "!": "potion",
+        "?": "scroll",
+        "/": "wand",
+        "=": "ring",
+        '"': "amulet",
+        "[": "armor",
+        ")": "weapon",
+        "(": "tool",
+        "%": "food",
+        "^": "trap",
+    }
+
+    # Monster symbols (letters)
+    for y, line in enumerate(lines):
+        for x, char in enumerate(line):
+            pos = (x, y)
+
+            if char == "@":
+                map_info["player_position"] = pos
+            elif char == "<":
+                map_info["stairs_positions"].append(("up", pos))
+            elif char == ">":
+                map_info["stairs_positions"].append(("down", pos))
+            elif char in ["+", "-", "|"]:
+                map_info["door_positions"].append((char, pos))
+            elif char in ["$", "*", "!", "?", "/", "=", '"', "[", ")", "(", "%"]:
+                map_info["item_positions"].append((char, pos))
+            elif char.isalpha() and char != "@":
+                # Store both the character and position for monster identification
+                map_info["monster_positions"].append((char, pos))
+            elif char == ".":
+                map_info["floor_positions"].append(pos)
+            elif char in ["#", " "]:
+                map_info["wall_positions"].append(pos)
+
+    return map_info
+
+
+def extract_game_context(observation: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extract high-level game context from observation.
+
+    Args:
+        observation: Raw observation dictionary
+
+    Returns:
+        Dictionary with game context information
+    """
+    context = {
+        "in_combat": False,
+        "in_shop": False,
+        "at_stairs": False,
+        "items_nearby": False,
+        "doors_nearby": False,
+        "low_health": False,
+        "hungry": False,
+        "encumbered": False,
+        "in_menu": observation.get("in_menu", False),
+        "game_over": observation.get("terminated", False),
+    }
+
+    # Parse map for context
+    if "ascii_map" in observation:
+        map_info = parse_ascii_map(observation["ascii_map"])
+
+        # Check if player is near important features
+        if map_info["player_position"]:
+            px, py = map_info["player_position"]
+
+            # Check for nearby monsters (within 2 squares)
+            # Exclude pets from combat detection
+            pet_symbols = ["f", "d"]  # f = kitten, d = dog
+            for monster_char, (mx, my) in map_info["monster_positions"]:
+                if abs(mx - px) <= 2 and abs(my - py) <= 2:
+                    # Only trigger combat for non-pet monsters
+                    if monster_char not in pet_symbols:
+                        context["in_combat"] = True
+                        break
+
+            # Check for stairs
+            for stair_type, (sx, sy) in map_info["stairs_positions"]:
+                if sx == px and sy == py:
+                    context["at_stairs"] = True
+                    context["stairs_type"] = stair_type
+                    break
+
+            # Check for nearby items
+            for _, (ix, iy) in map_info["item_positions"]:
+                if abs(ix - px) <= 1 and abs(iy - py) <= 1:
+                    context["items_nearby"] = True
+                    break
+
+            # Check for nearby doors
+            for _, (dx, dy) in map_info["door_positions"]:
+                if abs(dx - px) <= 1 and abs(dy - py) <= 1:
+                    context["doors_nearby"] = True
+                    break
+
+    # Check health status
+    if "character_stats" in observation:
+        stats = observation["character_stats"]
+        hp = stats.get("hp", 0)
+        max_hp = stats.get("max_hp", 1)
+        if hp < max_hp * 0.3:
+            context["low_health"] = True
+
+    # Check for specific messages
+    message = observation.get("message", "").lower()
+    if "hungry" in message or "weak" in message:
+        context["hungry"] = True
+    if "burdened" in message or "stressed" in message:
+        context["encumbered"] = True
+    if "shop" in message or "shopkeeper" in message:
+        context["in_shop"] = True
+
+    return context
+
+
+def simplify_observation(observation: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Create a simplified observation for agents that need less detail.
+
+    Args:
+        observation: Full observation dictionary
+
+    Returns:
+        Simplified observation dictionary
+    """
+    simplified = {
+        "turn": observation.get("turn_count", 0),
+        "level": observation.get("dungeon_level", 1),
+        "hp": 0,
+        "max_hp": 0,
+        "message": observation.get("message", ""),
+        "terminated": observation.get("terminated", False),
+        "reward": observation.get("reward_last", 0.0),
+    }
+
+    # Extract HP
+    if "character_stats" in observation:
+        stats = observation["character_stats"]
+        simplified["hp"] = stats.get("hp", 0)
+        simplified["max_hp"] = stats.get("max_hp", 0)
+
+    # Extract key map features
+    if "ascii_map" in observation:
+        map_info = parse_ascii_map(observation["ascii_map"])
+        simplified["player_pos"] = map_info["player_position"]
+        simplified["monsters_nearby"] = len(map_info["monster_positions"])
+        simplified["items_nearby"] = len(map_info["item_positions"])
+        simplified["at_stairs"] = any(
+            pos == map_info["player_position"] for _, pos in map_info["stairs_positions"]
+        )
+
+    return simplified
+
+
+def extract_inventory_from_message(message: str) -> List[Dict[str, Any]]:
+    """
+    Extract inventory information from NetHack inventory messages.
+
+    Args:
+        message: Inventory message string
+
+    Returns:
+        List of inventory items
+    """
+    items = []
+
+    # Common inventory line patterns
+    # Example: "a - a blessed +1 long sword (weapon in hand)"
+    # Example: "b - an uncursed food ration"
+    pattern = r"^([a-zA-Z])\s*-\s*(.+?)(?:\s*\(([^)]+)\))?$"
+
+    for line in message.split("\n"):
+        match = re.match(pattern, line.strip())
+        if match:
+            letter, description, status = match.groups()
+
+            item = {
+                "letter": letter,
+                "description": description.strip(),
+                "status": status.strip() if status else None,
+            }
+
+            # Parse quantity
+            qty_match = re.match(r"^(\d+)\s+(.+)", description)
+            if qty_match:
+                item["quantity"] = int(qty_match.group(1))
+                item["name"] = qty_match.group(2)
+            else:
+                item["quantity"] = 1
+                item["name"] = description
+
+            # Identify item type
+            item["type"] = identify_item_type(description)
+
+            items.append(item)
+
+    return items
+
+
+def identify_item_type(description: str) -> str:
+    """
+    Identify the type of an item from its description.
+
+    Args:
+        description: Item description string
+
+    Returns:
+        Item type string
+    """
+    desc_lower = description.lower()
+
+    # Weapons
+    if any(
+        word in desc_lower
+        for word in [
+            "sword",
+            "dagger",
+            "spear",
+            "axe",
+            "mace",
+            "bow",
+            "arrow",
+            "dart",
+            "knife",
+        ]
+    ):
+        return "weapon"
+
+    # Armor
+    if any(
+        word in desc_lower
+        for word in [
+            "armor",
+            "mail",
+            "helmet",
+            "boots",
+            "gloves",
+            "shield",
+            "cloak",
+            "robe",
+        ]
+    ):
+        return "armor"
+
+    # Food
+    if any(
+        word in desc_lower
+        for word in [
+            "food",
+            "ration",
+            "corpse",
+            "egg",
+            "fruit",
+            "meat",
+            "candy",
+            "cookie",
+        ]
+    ):
+        return "food"
+
+    # Potions
+    if "potion" in desc_lower:
+        return "potion"
+
+    # Scrolls
+    if "scroll" in desc_lower:
+        return "scroll"
+
+    # Wands
+    if "wand" in desc_lower:
+        return "wand"
+
+    # Rings
+    if "ring" in desc_lower:
+        return "ring"
+
+    # Amulets
+    if "amulet" in desc_lower:
+        return "amulet"
+
+    # Tools
+    if any(
+        word in desc_lower
+        for word in [
+            "pick",
+            "key",
+            "lamp",
+            "candle",
+            "bag",
+            "sack",
+            "horn",
+            "whistle",
+            "mirror",
+        ]
+    ):
+        return "tool"
+
+    # Gems/stones
+    if any(word in desc_lower for word in ["gem", "stone", "rock", "crystal"]):
+        return "gem"
+
+    # Gold
+    if "gold" in desc_lower:
+        return "gold"
+
+    return "unknown"

synth_ai/environments/examples/nethack/helpers/recording_wrapper.py

@@ -0,0 +1,201 @@
+"""Environment wrapper that adds trajectory recording capabilities."""
+
+from typing import Dict, Any, Optional, Tuple
+from pathlib import Path
+import logging
+
+from src.synth_env.examples.nethack.environment import NetHackEnvironment
+from src.synth_env.examples.nethack.helpers.trajectory_recorder import (
+    TrajectoryRecorder,
+)
+from src.synth_env.environment.tools import EnvToolCall
+
+
+logger = logging.getLogger(__name__)
+
+
+class RecordingNetHackEnvironment(NetHackEnvironment):
+    """NetHack environment with automatic trajectory recording."""
+
+    def __init__(
+        self,
+        save_dir: str = "temp/nethack_trajectories",
+        auto_record: bool = True,
+        **kwargs,
+    ):
+        """Initialize recording wrapper.
+
+        Args:
+            save_dir: Directory to save trajectories
+            auto_record: Whether to automatically record all episodes
+            **kwargs: Arguments passed to NetHackEnvironment
+        """
+        super().__init__(**kwargs)
+
+        self.save_dir = Path(save_dir)
+        self.save_dir.mkdir(parents=True, exist_ok=True)
+
+        self.auto_record = auto_record
+        self.recorder = TrajectoryRecorder(save_dir)
+        self.is_recording = False
+        self.trajectory_id = None
+
+    async def start(self, **kwargs) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+        """Start environment and optionally begin recording."""
+        public_state, private_state = await super().start(**kwargs)
+
+        if self.auto_record:
+            # Extract character role from task metadata
+            character_role = "adventurer"  # default
+            if hasattr(self.task_instance, "metadata") and hasattr(
+                self.task_instance.metadata, "character_role"
+            ):
+                character_role = self.task_instance.metadata.character_role
+
+            task_id = self.task_instance.task_id if self.task_instance else None
+            self.trajectory_id = self.recorder.start_recording(character_role, task_id)
+            self.is_recording = True
+
+            logger.info(f"Started recording trajectory: {self.trajectory_id}")
+
+            # Record initial state
+            obs = self._extract_observation(public_state)
+            self.recorder.record_step("reset", obs, 0.0, False, {})
+
+        return public_state, private_state
+
+    async def process_action(
+        self, tool_calls: list[EnvToolCall]
+    ) -> Tuple[Any, float, bool, Dict[str, Any]]:
+        """Process action and record if enabled."""
+        # Execute action
+        observation, reward, terminated, info = await super().process_action(tool_calls)
+
+        # Record step if recording
+        if self.is_recording and self.recorder.is_recording:
+            # Extract action from tool calls
+            action = "unknown"
+            if tool_calls and len(tool_calls) > 0:
+                if hasattr(tool_calls[0], "args") and "action" in tool_calls[0].args:
+                    action = tool_calls[0].args["action"]
+
+            # Extract observation data
+            obs_data = self._extract_observation(observation)
+
+            # Record step
+            self.recorder.record_step(action, obs_data, reward, terminated, info)
+
+        # If episode ended, stop recording
+        if terminated and self.is_recording:
+            await self._finalize_recording(observation)
+
+        return observation, reward, terminated, info
+
+    def _extract_observation(self, state: Any) -> Dict[str, Any]:
+        """Extract observation data from state object."""
+        if hasattr(state, "__dict__"):
+            # Convert state object to dict
+            obs = {}
+
+            # Extract relevant fields
+            if hasattr(state, "ascii_map"):
+                obs["ascii_map"] = state.ascii_map
+            if hasattr(state, "message"):
+                obs["message"] = state.message
+            if hasattr(state, "character_stats"):
+                obs["player_stats"] = state.character_stats
+            if hasattr(state, "inventory"):
+                obs["inventory"] = state.inventory
+            if hasattr(state, "position"):
+                obs["player_stats"] = obs.get("player_stats", {})
+                obs["player_stats"]["x"] = state.position[0]
+                obs["player_stats"]["y"] = state.position[1]
+            if hasattr(state, "dungeon_level"):
+                obs["player_stats"] = obs.get("player_stats", {})
+                obs["player_stats"]["depth"] = state.dungeon_level
+            if hasattr(state, "in_menu"):
+                obs["in_menu"] = state.in_menu
+            if hasattr(state, "menu_items"):
+                obs["menu_items"] = state.menu_items
+
+            return obs
+
+        # If it's already a dict, return as is
+        return state if isinstance(state, dict) else {}
+
+    async def _finalize_recording(self, final_state: Any):
+        """Finalize and save recording."""
+        if not self.is_recording:
+            return
+
+        # Determine final status
+        final_status = "completed"
+        if hasattr(final_state, "message"):
+            msg = final_state.message.lower()
+            if "die" in msg or "killed" in msg:
+                final_status = "died"
+            elif "quit" in msg:
+                final_status = "quit"
+            elif "time limit" in msg or "truncated" in msg:
+                final_status = "truncated"
+
+        # Extract achievements if available
+        achievements = {}
+        # TODO: Extract achievements from game state
+
+        # Stop recording
+        self.recorder.stop_recording(final_status, achievements)
+
+        # Save trajectory
+        filepath = self.recorder.save_trajectory()
+        logger.info(f"Saved trajectory to: {filepath}")
+
+        # Get and log summary
+        summary = self.recorder.get_summary()
+        logger.info(f"Trajectory summary: {summary}")
+
+        self.is_recording = False
+
+    def start_recording(
+        self, character_role: Optional[str] = None, task_id: Optional[str] = None
+    ) -> str:
+        """Manually start recording."""
+        if self.is_recording:
+            logger.warning("Recording already in progress")
+            return self.trajectory_id
+
+        if character_role is None:
+            character_role = "adventurer"
+            if hasattr(self, "engine") and hasattr(self.engine, "character_role"):
+                character_role = self.engine.character_role
+
+        self.trajectory_id = self.recorder.start_recording(character_role, task_id)
+        self.is_recording = True
+        logger.info(f"Started manual recording: {self.trajectory_id}")
+
+        return self.trajectory_id
+
+    def stop_recording(self, save: bool = True) -> Optional[str]:
+        """Manually stop recording."""
+        if not self.is_recording:
+            logger.warning("No recording in progress")
+            return None
+
+        self.recorder.stop_recording()
+
+        filepath = None
+        if save:
+            filepath = self.recorder.save_trajectory()
+            logger.info(f"Saved trajectory to: {filepath}")
+
+        self.is_recording = False
+        return filepath
+
+    def get_recording_status(self) -> Dict[str, Any]:
+        """Get current recording status."""
+        return {
+            "is_recording": self.is_recording,
+            "trajectory_id": self.trajectory_id,
+            "current_step": self.recorder.current_step if self.is_recording else 0,
+            "total_reward": self.recorder.total_reward if self.is_recording else 0.0,
+        }