synth-ai 0.2.9.dev11__py3-none-any.whl → 0.4.1__py3-none-any.whl

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

@@ -1,402 +0,0 @@
-"""NLE integration wrapper for NetHack environment."""
-
-from typing import Any, Dict, List, Optional, Tuple
-
-import numpy as np
-
-try:
-    import nle
-    from nle import nethack
-except ImportError as e:
-    raise ImportError(
-        "NLE (NetHack Learning Environment) is required but not installed. "
-        "Please install it with: pip install nle"
-    ) from e
-
-
-class NLEWrapper:
-    """Wrapper around NLE (NetHack Learning Environment) for synth-env integration."""
-
-    def __init__(self, character_role: str = "mon", observation_keys: Optional[List[str]] = None):
-        """Initialize NLE wrapper.
-
-        Args:
-            character_role: Character class (mon, val, wiz, etc.)
-            observation_keys: Which observations to include
-        """
-        self.character_role = self._convert_role_name(character_role)
-
-        # Default observation keys
-        if observation_keys is None:
-            observation_keys = [
-                "glyphs",
-                "chars",
-                "colors",
-                "specials",
-                "blstats",
-                "message",
-                "inv_glyphs",
-                "inv_letters",
-                "inv_oclasses",
-                "inv_strs",
-                "tty_chars",
-                "tty_colors",
-                "tty_cursor",
-            ]
-
-        # Create NLE environment
-        self.env = nle.env.NLE(character=self.character_role, observation_keys=observation_keys)
-
-        # Build action mapping
-        self._build_action_mapping()
-
-        # Track last observation
-        self.last_obs = None
-        self.last_reward = 0.0
-        self.last_done = False
-        self.last_info = {}
-
-    def _convert_role_name(self, role: str) -> str:
-        """Convert full role names to NLE abbreviations."""
-        role_map = {
-            "tourist": "tou",
-            "knight": "kni",
-            "wizard": "wiz",
-            "barbarian": "bar",
-            "ranger": "ran",
-            "priest": "pri",
-            "monk": "mon",
-            "rogue": "rog",
-            "valkyrie": "val",
-            "samurai": "sam",
-            "archeologist": "arc",
-            "healer": "hea",
-            "caveman": "cav",
-        }
-        return role_map.get(role.lower(), "mon")  # Default to monk
-
-    def _build_action_mapping(self):
-        """Build mapping from action names to NLE action indices."""
-        self.action_map = {}
-        self.index_to_action = {}
-
-        # Map each action in env.actions to a name
-        for i, action in enumerate(self.env.actions):
-            # Compass directions
-            if action == nethack.CompassDirection.N:
-                name = "north"
-            elif action == nethack.CompassDirection.S:
-                name = "south"
-            elif action == nethack.CompassDirection.E:
-                name = "east"
-            elif action == nethack.CompassDirection.W:
-                name = "west"
-            elif action == nethack.CompassDirection.NE:
-                name = "northeast"
-            elif action == nethack.CompassDirection.NW:
-                name = "northwest"
-            elif action == nethack.CompassDirection.SE:
-                name = "southeast"
-            elif action == nethack.CompassDirection.SW:
-                name = "southwest"
-
-            # Misc directions
-            elif action == nethack.MiscDirection.UP:
-                name = "up"
-            elif action == nethack.MiscDirection.DOWN:
-                name = "down"
-            elif action == nethack.MiscDirection.WAIT:
-                name = "wait"
-
-            # Commands
-            elif action == nethack.Command.SEARCH:
-                name = "search"
-            elif action == nethack.Command.INVENTORY:
-                name = "inventory"
-            elif action == nethack.Command.LOOK:
-                name = "look"
-            elif action == nethack.Command.OPEN:
-                name = "open"
-            elif action == nethack.Command.CLOSE:
-                name = "close"
-            elif action == nethack.Command.KICK:
-                name = "kick"
-            elif action == nethack.Command.PICKUP:
-                name = "pickup"
-            elif action == nethack.Command.DROP:
-                name = "drop"
-            elif action == nethack.Command.EAT:
-                name = "eat"
-            elif action == nethack.Command.WIELD:
-                name = "wield"
-            elif action == nethack.Command.WEAR:
-                name = "wear"
-            elif action == nethack.Command.TAKEOFF:
-                name = "takeoff"
-            elif action == nethack.Command.PUTON:
-                name = "puton"
-            elif action == nethack.Command.REMOVE:
-                name = "remove"
-            elif action == nethack.Command.READ:
-                name = "read"
-            elif action == nethack.Command.QUAFF:
-                name = "quaff"
-            elif action == nethack.Command.ZAP:
-                name = "zap"
-            elif action == nethack.Command.THROW:
-                name = "throw"
-            elif action == nethack.Command.FIRE:
-                name = "fire"
-            elif action == nethack.Command.APPLY:
-                name = "apply"
-            elif action == nethack.Command.PRAY:
-                name = "pray"
-            elif action == nethack.Command.SAVE:
-                name = "save"
-            elif action == nethack.Command.QUIT:
-                name = "quit"
-            elif action == nethack.Command.ESC:
-                name = "escape"
-            elif action == nethack.Command.PAY:
-                name = "pay"
-            elif action == nethack.Command.LOOT:
-                name = "loot"
-            elif action == nethack.Command.ENHANCE:
-                name = "enhance"
-            elif action == nethack.Command.FORCE:
-                name = "force"
-            elif action == nethack.Command.INVOKE:
-                name = "invoke"
-            elif action == nethack.Command.OFFER:
-                name = "offer"
-            elif action == nethack.Command.RUB:
-                name = "rub"
-            elif action == nethack.Command.SIT:
-                name = "sit"
-            elif action == nethack.Command.TURN:
-                name = "turn"
-            elif action == nethack.Command.UNTRAP:
-                name = "untrap"
-            elif action == nethack.Command.WIPE:
-                name = "wipe"
-            elif action == nethack.Command.ENGRAVE:
-                name = "engrave"
-            elif action == nethack.Command.JUMP:
-                name = "jump"
-            elif action == nethack.Command.CHAT:
-                name = "chat"
-            elif action == nethack.Command.DIP:
-                name = "dip"
-            elif action == nethack.Command.RIDE:
-                name = "ride"
-            elif action == nethack.Command.TIP:
-                name = "tip"
-
-            # Special/Misc
-            elif action == nethack.MiscAction.MORE:
-                name = "more"
-
-            # Text characters (for menu selection)
-            elif hasattr(action, "value") and 97 <= action.value <= 122:  # a-z
-                name = chr(action.value)
-            elif hasattr(action, "value") and 65 <= action.value <= 90:  # A-Z
-                name = chr(action.value)
-            elif hasattr(action, "value") and 48 <= action.value <= 57:  # 0-9
-                name = chr(action.value)
-            elif action == nethack.TextCharacters.SPACE:
-                name = "space"
-            elif action == nethack.TextCharacters.APOS:
-                name = "'"
-            elif action == nethack.TextCharacters.QUOTE:
-                name = '"'
-            else:
-                # Skip unmapped actions
-                continue
-
-            self.action_map[name] = i
-            self.index_to_action[i] = name
-
-    def reset(self, seed: Optional[int] = None) -> Dict[str, Any]:
-        """Reset the NLE environment."""
-        if seed is not None:
-            self.env.seed(seed)
-
-        self.last_obs = self.env.reset()
-        self.last_reward = 0.0
-        self.last_done = False
-        self.last_info = {}
-
-        return self._process_observation(self.last_obs)
-
-    def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
-        """Take a step in the environment.
-
-        Args:
-            action: Action name (e.g., "north", "pickup", "a" for menu)
-
-        Returns:
-            Tuple of (observation, reward, done, info)
-        """
-        # Handle compound actions like "open west"
-        if " " in action:
-            # Split compound action
-            parts = action.split(" ", 1)
-            if len(parts) == 2 and parts[0] in self.action_map and parts[1] in self.action_map:
-                # Execute both actions in sequence
-                print(f"Splitting compound action '{action}' into '{parts[0]}' then '{parts[1]}'")
-                # First action
-                action_idx = self.action_map[parts[0]]
-                self.env.step(action_idx)
-                # Second action (the direction)
-                action = parts[1]
-            else:
-                raise ValueError(f"Invalid compound action: {action}. Use separate actions.")
-
-        # Get action index
-        if action not in self.action_map:
-            # Special handling for menu letters that might not be in action_map
-            if len(action) == 1 and action.isalpha():
-                # Try to find ASCII value for single letter
-                ascii_val = ord(action)
-                # Look for an action with this ASCII value
-                for i, act in enumerate(self.env.actions):
-                    if hasattr(act, "value") and act.value == ascii_val:
-                        print(f"Found menu letter '{action}' at action index {i}")
-                        action_idx = i
-                        break
-                else:
-                    raise ValueError(
-                        f"Unknown action: {action}. Valid actions: {list(self.action_map.keys())}"
-                    )
-            else:
-                raise ValueError(
-                    f"Unknown action: {action}. Valid actions: {list(self.action_map.keys())}"
-                )
-        else:
-            action_idx = self.action_map[action]
-
-        # Take step
-        self.last_obs, self.last_reward, self.last_done, self.last_info = self.env.step(action_idx)
-
-        # Process observation
-        processed_obs = self._process_observation(self.last_obs)
-
-        return processed_obs, self.last_reward, self.last_done, self.last_info
-
-    def _process_observation(self, obs: Dict[str, np.ndarray]) -> Dict[str, Any]:
-        """Process NLE observation into a more usable format."""
-        processed = {}
-
-        # Decode message
-        if "message" in obs:
-            processed["message"] = obs["message"].tobytes().decode("ascii").strip()
-
-        # Get ASCII map
-        if "chars" in obs:
-            processed["ascii_chars"] = obs["chars"]
-            # Convert to string map
-            lines = []
-            for row in obs["chars"]:
-                line = "".join(chr(c) for c in row)
-                lines.append(line)
-            processed["ascii_map"] = "\n".join(lines)
-
-        # Get player stats from blstats
-        if "blstats" in obs:
-            blstats = obs["blstats"]
-            processed["player_stats"] = {
-                "x": int(blstats[0]),
-                "y": int(blstats[1]),
-                "strength": int(blstats[2]),
-                "strength_pct": int(blstats[3]),
-                "dexterity": int(blstats[4]),
-                "constitution": int(blstats[5]),
-                "intelligence": int(blstats[6]),
-                "wisdom": int(blstats[7]),
-                "charisma": int(blstats[8]),
-                "score": int(blstats[9]),
-                "hp": int(blstats[10]),
-                "max_hp": int(blstats[11]),
-                "depth": int(blstats[12]),
-                "gold": int(blstats[13]),
-                "energy": int(blstats[14]),
-                "max_energy": int(blstats[15]),
-                "ac": int(blstats[16]),
-                "monster_level": int(blstats[17]),
-                "experience_level": int(blstats[18]),
-                "experience_points": int(blstats[19]),
-                "time": int(blstats[20]),
-                "hunger_state": int(blstats[21]),
-                "carrying_capacity": int(blstats[22]),
-                "dungeon_number": int(blstats[23]),
-                "level_number": int(blstats[24]),
-            }
-
-        # Get inventory
-        if "inv_strs" in obs:
-            inv_items = []
-            for i, inv_str in enumerate(obs["inv_strs"]):
-                if inv_str[0] != 0:  # Non-empty slot
-                    item_str = inv_str.tobytes().decode("ascii").strip("\x00")
-                    if item_str:
-                        letter = chr(obs["inv_letters"][i]) if "inv_letters" in obs else "?"
-                        inv_items.append({"letter": letter, "description": item_str})
-            processed["inventory"] = inv_items
-
-        # Check if in menu
-        if "tty_chars" in obs:
-            tty_text = []
-            for row in obs["tty_chars"][:5]:  # Check first 5 rows
-                line = "".join(chr(c) if 32 <= c <= 126 else " " for c in row).strip()
-                if line:
-                    tty_text.append(line)
-
-            # Simple menu detection
-            processed["in_menu"] = any(
-                keyword in " ".join(tty_text).lower()
-                for keyword in [
-                    "pick up",
-                    "drop",
-                    "wear",
-                    "take off",
-                    "what do you want",
-                ]
-            )
-            processed["menu_text"] = tty_text
-
-        # TTY cursor position
-        if "tty_cursor" in obs:
-            processed["cursor"] = (int(obs["tty_cursor"][1]), int(obs["tty_cursor"][0]))
-
-        # Raw observations for advanced use
-        processed["_raw"] = obs
-
-        return processed
-
-    def get_valid_actions(self) -> List[str]:
-        """Get list of valid action names."""
-        return list(self.action_map.keys())
-
-    def get_state(self) -> bytes:
-        """Get the current NLE state for serialization."""
-        # NLE environments have a clone_state method that returns the full state
-        if hasattr(self.env, "clone_state"):
-            return self.env.clone_state()
-        elif hasattr(self.env, "clone_full_state"):
-            return self.env.clone_full_state()
-        else:
-            raise RuntimeError("NLE environment does not support state cloning")
-
-    def set_state(self, state: bytes):
-        """Restore NLE state from serialized data."""
-        # NLE environments have a restore_state method
-        if hasattr(self.env, "restore_state"):
-            self.env.restore_state(state)
-        elif hasattr(self.env, "restore_full_state"):
-            self.env.restore_full_state(state)
-        else:
-            raise RuntimeError("NLE environment does not support state restoration")
-
-    def close(self):
-        """Close the NLE environment."""
-        self.env.close()