synth-ai 0.2.0__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

synth_ai/environments/examples/nethack/environment.py

@@ -0,0 +1,255 @@
+"""NetHack environment wrapper for synth-env framework."""
+
+from __future__ import annotations
+
+import json
+from typing import Dict, Any, Optional, List, Union
+from pydantic import BaseModel
+
+from synth_ai.environments.stateful.core import StatefulEnvironment
+from synth_ai.environments.reproducibility.core import ReproducibleEnvironment
+from synth_ai.environments.environment.tools import AbstractTool, EnvToolCall, ToolResult
+from synth_ai.environments.environment.shared_engine import (
+    GetObservationCallable,
+    InternalObservation,
+)
+from synth_ai.environments.tasks.core import TaskInstance
+
+from .engine import (
+    NetHackEngine,
+    NetHackObservationCallable,
+    NetHackCheckpointObservationCallable,
+    NetHackPublicState,
+    NetHackPrivateState,
+)
+from .helpers import (
+    validate_action,
+    get_action_description,
+    NETHACK_ACTIONS,
+    MENU_ACTIONS,
+)
+
+
+class NetHackActionInput(BaseModel):
+    """Pydantic model for NetHack action validation."""
+
+    action: str  # Action string from NETHACK_ACTIONS or MENU_ACTIONS
+
+
+class NetHackInteractTool(AbstractTool):
+    """Tool for performing actions in NetHack."""
+
+    name = "interact"
+    description = (
+        "Perform an action in the NetHack dungeon. Available actions include "
+        "movement (north, south, east, west), combat (fight), inventory management "
+        "(inventory, pickup, drop), and many others. In menus, use letter keys (a-z) "
+        "or numbers (0-9) to select options, or 'escape' to cancel."
+    )
+    call_schema = NetHackActionInput
+    result_schema = ToolResult
+
+    def __init__(self, engine: NetHackEngine):
+        self.engine = engine
+
+    async def __call__(self, call: EnvToolCall) -> ToolResult:
+        """Execute the interact tool."""
+        try:
+            action = call.args["action"]  # Will KeyError if missing
+
+            # Get current game state for validation
+            priv, pub = self.engine.get_current_states_for_observation()
+            game_state = {
+                "in_menu": pub.in_menu,
+                "terminated": pub.terminated,
+                "stairs_here": False,  # Would be determined from map parsing
+            }
+
+            # Validate action
+            is_valid, error_msg = validate_action(action, game_state)
+            if not is_valid:
+                return ToolResult(
+                    ok=False,
+                    error=error_msg or f"Invalid action: {action}",
+                    payload={"public_state": pub, "private_state": priv},
+                )
+
+            # Execute action
+            private_state, public_state = await self.engine._step_engine(action)
+
+            return ToolResult(
+                ok=True,
+                payload={"public_state": public_state, "private_state": private_state},
+            )
+
+        except Exception as e:
+            # Return current state even on error
+            priv, pub = self.engine.get_current_states_for_observation()
+            return ToolResult(
+                ok=False,
+                error=str(e),
+                payload={"public_state": pub, "private_state": priv},
+            )
+
+
+class NetHackEnvironment(StatefulEnvironment, ReproducibleEnvironment[NetHackEngine]):
+    """NetHack environment implementation."""
+
+    def __init__(
+        self,
+        task_instance: TaskInstance,
+        custom_step_obs: Optional[GetObservationCallable] = None,
+        custom_ckpt_obs: Optional[GetObservationCallable] = None,
+    ):
+        """Initialize NetHack environment."""
+        self.name = "NetHack"
+        self.task_instance = task_instance
+        self.custom_step_observation_callable = custom_step_obs or NetHackObservationCallable()
+        self.custom_checkpoint_observation_callable = (
+            custom_ckpt_obs or NetHackCheckpointObservationCallable()
+        )
+        self.engine = NetHackEngine(task_instance)
+        self._interact_tool = NetHackInteractTool(self.engine)
+
+    async def initialize(self) -> InternalObservation:
+        """Initialize the environment and return initial observation."""
+        priv, pub = await self.engine._reset_engine()
+        return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+
+    async def step(
+        self, tool_calls: Union[List[EnvToolCall], EnvToolCall, Dict, List[Dict], str]
+    ) -> InternalObservation:
+        """Execute one step in the environment."""
+        try:
+            validated_call = self.validate_tool_calls(tool_calls)
+        except ValueError as e:
+            # Return current state with error
+            priv, pub = self.engine.get_current_states_for_observation()
+            return await self._to_observation(
+                priv,
+                pub,
+                self.custom_step_observation_callable,
+                extra_obs={"error": str(e)},
+            )
+
+        # Execute the tool
+        result = await self._interact_tool(validated_call)
+
+        if result.ok:
+            priv = result.payload["private_state"]
+            pub = result.payload["public_state"]
+            return await self._to_observation(priv, pub, self.custom_step_observation_callable)
+        else:
+            # Tool failed - return error with current state
+            priv, pub = self.engine.get_current_states_for_observation()
+            return await self._to_observation(
+                priv,
+                pub,
+                self.custom_step_observation_callable,
+                extra_obs={"error": result.error},
+            )
+
+    async def checkpoint(self) -> InternalObservation:
+        """Create a checkpoint observation."""
+        priv, pub = self.engine.get_current_states_for_observation()
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    async def terminate(self) -> InternalObservation:
+        """Terminate the environment."""
+        priv, pub = self.engine.get_current_states_for_observation()
+
+        # Mark as terminated
+        pub.terminated = True
+        priv.terminated = True
+
+        return await self._to_observation(priv, pub, self.custom_checkpoint_observation_callable)
+
+    def validate_tool_calls(
+        self, tool_calls: Union[List[EnvToolCall], EnvToolCall, Dict, List[Dict], str]
+    ) -> EnvToolCall:
+        """Validate and normalize tool calls."""
+        # Handle string input (simple action)
+        if isinstance(tool_calls, str):
+            return EnvToolCall(tool="interact", args={"action": tool_calls})
+
+        # Handle dict input
+        if isinstance(tool_calls, dict):
+            # Check if it's already properly formatted
+            if "tool" in tool_calls and "args" in tool_calls:
+                # Handle tool name aliases
+                tool_name = tool_calls["tool"]
+                if tool_name == "nethack_interact":
+                    tool_name = "interact"
+                return EnvToolCall(tool=tool_name, args=tool_calls["args"])  # type: ignore[misc]
+            elif "tool_name" in tool_calls and "args" in tool_calls:
+                # Handle legacy format
+                tool_name = tool_calls["tool_name"]
+                if tool_name == "nethack_interact":
+                    tool_name = "interact"
+                return EnvToolCall(tool=tool_name, args=tool_calls["args"])
+            # Check for action key
+            elif "action" in tool_calls:
+                return EnvToolCall(tool="interact", args={"action": tool_calls["action"]})
+            # Check for tool_calls format
+            elif "tool_calls" in tool_calls:
+                tool_calls = tool_calls["tool_calls"]
+                if isinstance(tool_calls, list) and len(tool_calls) > 0:
+                    return self.validate_tool_calls(tool_calls[0])
+            # Try to extract action from various formats
+            else:
+                # Look for action in nested structures
+                for key in ["args", "parameters", "input"]:
+                    if key in tool_calls and isinstance(tool_calls[key], dict):
+                        if "action" in tool_calls[key]:
+                            return EnvToolCall(
+                                tool="interact",
+                                args={"action": tool_calls[key]["action"]},
+                            )
+
+        # Handle list input
+        if isinstance(tool_calls, list):
+            if len(tool_calls) == 0:
+                raise ValueError("Empty tool calls list")
+            # Take first tool call
+            return self.validate_tool_calls(tool_calls[0])
+
+        # Handle EnvToolCall object
+        if isinstance(tool_calls, EnvToolCall):
+            return tool_calls
+
+        raise ValueError(
+            f"Invalid tool call format. Expected action string, dict with 'action' key, "
+            f"or EnvToolCall object. Got: {type(tool_calls)}"
+        )
+
+    async def _to_observation(
+        self,
+        private_state: NetHackPrivateState,
+        public_state: NetHackPublicState,
+        observation_callable: GetObservationCallable,
+        extra_obs: Optional[Dict[str, Any]] = None,
+    ) -> InternalObservation:
+        """Convert states to observation using the callable."""
+        obs = await observation_callable.get_observation(public_state, private_state)  # type: ignore[call-arg]
+
+        if extra_obs:
+            obs.update(extra_obs)
+
+        return obs
+
+    async def _serialize_engine(self) -> Any:
+        """Serialize the engine state."""
+        return await self.engine._serialize_engine()
+
+    @classmethod
+    async def _deserialize_engine(cls, task_instance: TaskInstance, snapshot: Any) -> NetHackEngine:
+        """Deserialize the engine from a snapshot."""
+        return await NetHackEngine._deserialize_engine(snapshot)
+
+    def get_available_actions(self) -> List[str]:
+        """Get list of all available actions."""
+        return list(NETHACK_ACTIONS.keys()) + list(MENU_ACTIONS.keys())
+
+    def get_action_descriptions(self) -> Dict[str, str]:
+        """Get descriptions of all actions."""
+        return {**NETHACK_ACTIONS, **MENU_ACTIONS}

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

@@ -0,0 +1,42 @@
+"""Helper utilities for NetHack environment."""
+
+from .action_mapping import (
+    NETHACK_ACTIONS,
+    MENU_ACTIONS,
+    ALL_ACTIONS,
+    ACTION_CATEGORIES,
+    validate_action,
+    get_action_description,
+    get_actions_for_context,
+    convert_action_to_nle,
+    parse_compound_action,
+)
+
+from .observation_utils import (
+    format_observation_for_llm,
+    parse_ascii_map,
+    extract_game_context,
+    simplify_observation,
+    extract_inventory_from_message,
+    identify_item_type,
+)
+
+__all__ = [
+    # Action mapping
+    "NETHACK_ACTIONS",
+    "MENU_ACTIONS",
+    "ALL_ACTIONS",
+    "ACTION_CATEGORIES",
+    "validate_action",
+    "get_action_description",
+    "get_actions_for_context",
+    "convert_action_to_nle",
+    "parse_compound_action",
+    # Observation utils
+    "format_observation_for_llm",
+    "parse_ascii_map",
+    "extract_game_context",
+    "simplify_observation",
+    "extract_inventory_from_message",
+    "identify_item_type",
+]

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

@@ -0,0 +1,301 @@
+"""Action mapping and validation for NetHack."""
+
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass
+
+
+@dataclass
+class ActionCategory:
+    """Category of actions with description."""
+
+    name: str
+    description: str
+    actions: List[str]
+
+
+# Comprehensive NetHack action mapping
+NETHACK_ACTIONS: Dict[str, str] = {
+    # Movement actions (8 directions + wait)
+    "north": "move north",
+    "south": "move south",
+    "east": "move east",
+    "west": "move west",
+    "northeast": "move northeast",
+    "northwest": "move northwest",
+    "southeast": "move southeast",
+    "southwest": "move southwest",
+    "wait": "wait/rest for one turn",
+    # Movement modifiers
+    "run_north": "run north until something interesting",
+    "run_south": "run south until something interesting",
+    "run_east": "run east until something interesting",
+    "run_west": "run west until something interesting",
+    "go_up": "go up stairs/ladder",
+    "go_down": "go down stairs/ladder",
+    # Basic interactions
+    "search": "search for secret doors/traps",
+    "open": "open a door",
+    "close": "close a door",
+    "kick": "kick something",
+    "force": "force a lock",
+    "untrap": "untrap something",
+    # Inventory and items
+    "inventory": "check inventory",
+    "pickup": "pick up items",
+    "drop": "drop items",
+    "dropall": "drop all items",
+    "wear": "wear armor/accessories",
+    "take_off": "take off armor/accessories",
+    "wield": "wield a weapon",
+    "unwield": "unwield current weapon",
+    "quiver": "ready ammunition",
+    "put_on": "put on accessories",
+    "remove": "remove accessories",
+    # Using items
+    "eat": "eat food",
+    "drink": "drink a potion",
+    "read": "read a scroll/spellbook",
+    "zap": "zap a wand",
+    "apply": "apply/use a tool",
+    "invoke": "invoke an artifact",
+    "rub": "rub a lamp/stone",
+    "throw": "throw an item",
+    "fire": "fire from quiver",
+    # Magic
+    "cast": "cast a spell",
+    "pray": "pray to your deity",
+    "offer": "offer sacrifice",
+    "turn_undead": "turn undead (priest ability)",
+    # Information (NOTE: These don't consume turns!)
+    "look": "look around (FREE ACTION - doesn't advance time)",
+    "farlook": "look at specific location (FREE ACTION)",
+    "whatis": "identify map symbol (FREE ACTION)",
+    "identify": "identify inventory item (FREE ACTION)",
+    "discoveries": "list discoveries (FREE ACTION)",
+    "conduct": "check conduct (FREE ACTION)",
+    "attributes": "check attributes (FREE ACTION)",
+    # Character actions
+    "enhance": "enhance skills",
+    "sit": "sit down",
+    "pay": "pay shopkeeper",
+    "chat": "talk to someone",
+    "loot": "loot a container",
+    "engrave": "write on the ground",
+    "monster_ability": "use monster ability",
+    # Game commands
+    "save": "save the game",
+    "quit": "quit the game",
+    "help": "show help",
+    "version": "show version",
+    "history": "show message history",
+    "name": "name an item/monster",
+    "call": "call item type",
+    "adjust": "adjust inventory letters",
+    # Special responses for prompts/menus
+    "yes": "answer yes",
+    "no": "answer no",
+    "all": "select all",
+    "none": "select none",
+    "menu_next": "next menu page",
+    "menu_previous": "previous menu page",
+    "escape": "cancel/escape",
+}
+
+# Single character responses for menu selections
+MENU_ACTIONS: Dict[str, str] = {
+    chr(i): f"select option {chr(i)}" for i in range(ord("a"), ord("z") + 1)
+}
+MENU_ACTIONS.update({chr(i): f"select option {chr(i)}" for i in range(ord("A"), ord("Z") + 1)})
+MENU_ACTIONS.update({str(i): f"select option {i}" for i in range(10)})
+
+# Combine all actions
+ALL_ACTIONS = {**NETHACK_ACTIONS, **MENU_ACTIONS}
+
+# Action categories for organization
+ACTION_CATEGORIES = [
+    ActionCategory(
+        name="Movement",
+        description="Basic movement and navigation",
+        actions=[
+            "north",
+            "south",
+            "east",
+            "west",
+            "northeast",
+            "northwest",
+            "southeast",
+            "southwest",
+            "wait",
+            "go_up",
+            "go_down",
+        ],
+    ),
+    ActionCategory(
+        name="Inventory",
+        description="Managing items and equipment",
+        actions=[
+            "inventory",
+            "pickup",
+            "drop",
+            "wear",
+            "wield",
+            "eat",
+            "drink",
+            "read",
+            "apply",
+            "throw",
+        ],
+    ),
+    ActionCategory(
+        name="Combat",
+        description="Fighting and defense (attack by moving into monsters!)",
+        actions=["fire", "zap", "throw", "kick"],
+    ),
+    ActionCategory(
+        name="Exploration",
+        description="Discovering the dungeon",
+        actions=["search", "open", "close", "look", "farlook"],
+    ),
+    ActionCategory(
+        name="Magic",
+        description="Spells and divine intervention",
+        actions=["cast", "pray", "offer", "invoke"],
+    ),
+    ActionCategory(
+        name="Game",
+        description="Meta game commands",
+        actions=["save", "quit", "help", "inventory"],
+    ),
+]
+
+
+def validate_action(action: str, game_state: Optional[Dict] = None) -> Tuple[bool, Optional[str]]:
+    """
+    Validate if an action is valid given the current game state.
+
+    Args:
+        action: The action string to validate
+        game_state: Optional game state for context-aware validation
+
+    Returns:
+        Tuple of (is_valid, error_message)
+    """
+    # Check if action exists
+    if action not in ALL_ACTIONS:
+        return False, f"Unknown action: {action}. Use 'help' to see available actions."
+
+    # Context-aware validation if game state provided
+    if game_state:
+        # Check if in menu
+        if game_state.get("in_menu", False):
+            if action not in MENU_ACTIONS and action not in [
+                "escape",
+                "menu_next",
+                "menu_previous",
+            ]:
+                return False, "Currently in a menu. Use letter/number keys or 'escape'."
+
+        # Check if game is over
+        if game_state.get("terminated", False):
+            if action not in ["quit", "save"]:
+                return False, "Game is over. You can only 'save' or 'quit'."
+
+        # Check stairs availability
+        if action in ["go_up", "go_down"]:
+            if not game_state.get("stairs_here", False):
+                return False, f"No stairs here to {action.replace('go_', '')}."
+
+    return True, None
+
+
+def get_action_description(action: str) -> str:
+    """Get the description of an action."""
+    return ALL_ACTIONS.get(action, "Unknown action")
+
+
+def get_actions_for_context(game_state: Dict) -> List[str]:
+    """Get relevant actions for the current game context."""
+    if game_state.get("in_menu", False):
+        # In menu - return menu navigation actions
+        menu_items = game_state.get("menu_items", [])
+        actions = ["escape"]
+
+        # Add letter options based on menu items
+        for i, item in enumerate(menu_items):
+            if i < 26:
+                actions.append(chr(ord("a") + i))
+
+        return actions
+
+    if game_state.get("terminated", False):
+        return ["quit", "save"]
+
+    # Normal gameplay - return common actions
+    common_actions = [
+        "north",
+        "south",
+        "east",
+        "west",
+        "search",
+        "inventory",
+        "pickup",
+        "look",
+        "wait",
+        "open",
+        "close",
+    ]
+
+    # Add context-specific actions
+    if game_state.get("stairs_here", False):
+        if game_state.get("stairs_down", False):
+            common_actions.append("go_down")
+        if game_state.get("stairs_up", False):
+            common_actions.append("go_up")
+
+    if game_state.get("items_here", False):
+        common_actions.append("pickup")
+
+    if game_state.get("door_here", False):
+        if game_state.get("door_open", False):
+            common_actions.append("close")
+        else:
+            common_actions.append("open")
+
+    return common_actions
+
+
+def convert_action_to_nle(action: str, action_map: Dict[str, int]) -> Optional[int]:
+    """
+    Convert string action to NLE integer action.
+
+    Args:
+        action: String action name
+        action_map: Dictionary mapping action names to NLE indices
+
+    Returns:
+        NLE action index or None if not found
+    """
+    # Direct lookup in action map
+    if action in action_map:
+        return action_map[action]
+
+    # Handle special cases
+    if action == "terminate":
+        # This is handled at a higher level
+        return None
+
+    # Single character actions (menu selections)
+    if len(action) == 1 and (action.isalpha() or action.isdigit()):
+        if action in action_map:
+            return action_map[action]
+
+    return None
+
+
+def parse_compound_action(action: str) -> List[str]:
+    """
+    Parse compound actions into individual steps.
+    E.g., "go to stairs and descend" -> ["navigate_to_stairs", "go_down"]
+    """
+    # This could be extended to handle more complex action parsing
+    return [action]  # For now, just return the action as-is