cogames-agents 0.0.0.7__cp312-cp312-macosx_11_0_arm64.whl

cogames_agents/policy/scripted_agent/unclipping_agent.py

@@ -0,0 +1,461 @@
+"""
+UnclippingAgent - Extends SimpleBaselineAgent with unclipping capabilities.
+
+This agent can detect clipped extractors and craft unclip items to restore them.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Optional, cast
+
+from mettagrid.policy.policy import MultiAgentPolicy, StatefulAgentPolicy
+from mettagrid.policy.policy_env_interface import PolicyEnvInterface
+from mettagrid.simulator import Action
+from mettagrid.simulator.interface import AgentObservation
+
+from .baseline_agent import BaselineAgentPolicyImpl
+from .types import (
+    BaselineHyperparameters,
+    ExtractorInfo,
+    Phase,
+    SimpleAgentState,
+)
+from .utils import is_adjacent, use_object_at
+
+if TYPE_CHECKING:
+    pass
+
+
+@dataclass
+class UnclippingHyperparameters(BaselineHyperparameters):
+    """Extends baseline hyperparameters with unclipping-specific parameters."""
+
+    # Unclipping strategy
+    unclip_priority_order: tuple[str, ...] = ("oxygen", "silicon", "carbon", "germanium")  # Order to unclip resources
+    craft_unclip_items_early: bool = True  # Craft unclip items proactively vs on-demand
+
+
+@dataclass
+class UnclippingAgentState(SimpleAgentState):
+    """Extended state for unclipping agent.
+
+    Note: decoder, modulator, resonator, scrambler are already defined in SimpleAgentState.
+    """
+
+    # Unclip tracking
+    blocked_by_clipped_extractor: Optional[tuple[int, int]] = None
+    unclip_target_resource: Optional[str] = None  # Which resource is clipped
+
+    # Discovered unclip recipes from observations (resource_type -> (unclip_item, craft_recipe))
+    # e.g., "oxygen" -> ("decoder", {"carbon": 1})
+    unclip_recipes: dict[str, tuple[str, dict[str, int]]] = field(default_factory=dict)
+
+    # Craft recipes for unclip items (item_name -> craft_recipe)
+    # e.g., "decoder" -> {"carbon": 1}
+    # Initialized from hub protocols
+    unclip_craft_recipes: dict[str, dict[str, int]] = field(default_factory=dict)
+
+
+class UnclippingAgentPolicyImpl(BaselineAgentPolicyImpl):
+    """
+    Agent that can unclip extractors by crafting and using unclip items.
+
+    Unclip item mapping:
+    - decoder (from carbon) unclips oxygen extractors
+    - modulator (from oxygen) unclips carbon extractors
+    - resonator (from silicon) unclips germanium extractors
+    - scrambler (from germanium) unclips silicon extractors
+    """
+
+    def __init__(
+        self,
+        policy_env_info: PolicyEnvInterface,
+        agent_id: int,
+        hyperparams: UnclippingHyperparameters,
+    ):
+        super().__init__(policy_env_info, agent_id, hyperparams)
+        self._priority_order = hyperparams.unclip_priority_order
+
+    def initial_agent_state(self) -> UnclippingAgentState:
+        """Create initial state for unclipping agent."""
+        # Get the base state from parent class
+        base_state = super().initial_agent_state()
+
+        # Initialize unclip recipes from hub protocols
+        # Map: unclip_item -> craft_recipe (e.g., "decoder" -> {"carbon": 1})
+        unclip_craft_recipes = {}
+        for protocol in getattr(self._policy_env_info, "hub_protocols", []):
+            for output_item, output_amount in protocol.output_resources.items():
+                if output_amount > 0 and output_item in ("decoder", "modulator", "resonator", "scrambler"):
+                    craft_recipe = dict(protocol.input_resources)
+                    craft_recipe.pop("energy", None)
+                    unclip_craft_recipes[output_item] = craft_recipe
+
+        # Create unclipping state by extending base state
+        # Convert base_state dict to UnclippingAgentState with additional fields
+        return UnclippingAgentState(
+            **base_state.__dict__,
+            unclip_recipes={},  # Will be discovered from clipped extractor observations
+            unclip_craft_recipes=unclip_craft_recipes,  # Store craft recipes separately
+        )
+
+    def _discover_unclip_recipe(self, s: UnclippingAgentState, parsed_observation) -> None:
+        """
+        Discover unclip recipes from clipped extractor observations.
+
+        When we observe a clipped extractor, its protocol_outputs tell us what unclip item is needed.
+        We then need to observe the hub with the correct vibe to learn how to craft that item.
+        """
+        for _pos, obj_state in parsed_observation.nearby_objects.items():
+            obj_name = obj_state.name.lower()
+            # Check if this is a clipped extractor
+            if "extractor" in obj_name and obj_state.clipped > 0:
+                # Extract resource type from name (e.g., "oxygen_extractor" -> "oxygen")
+                resource_type = obj_name.replace("_extractor", "")
+                # Check if we already know this recipe
+                if resource_type in s.unclip_recipes:
+                    continue
+
+                # Read the unclip item from protocol_inputs (unclipping protocols have inputs, not outputs!)
+                if obj_state.protocol_inputs:
+                    # The clipped extractor's protocol_inputs shows what unclip item is needed
+                    for item_name, amount in obj_state.protocol_inputs.items():
+                        if amount > 0 and item_name in ("decoder", "modulator", "resonator", "scrambler"):
+                            # Get the craft recipe from our initialized recipes
+                            craft_recipe = s.unclip_craft_recipes.get(item_name, {})
+                            s.unclip_recipes[resource_type] = (item_name, craft_recipe)
+                            break
+
+            # No need to discover craft recipes from hub - we already have them from init!
+
+    def _get_unclip_item_name(self, s: UnclippingAgentState, clipped_resource: str) -> Optional[str]:
+        recipe = s.unclip_recipes.get(clipped_resource)
+        if recipe is None:
+            return None
+        item_name, _ = recipe
+        return item_name
+
+    def _has_unclip_item(self, s: UnclippingAgentState) -> bool:
+        """Check if agent has the unclip item for the blocked extractor."""
+        if s.blocked_by_clipped_extractor is None:
+            return False
+
+        if s.unclip_target_resource is None:
+            return False
+
+        item_name = self._get_unclip_item_name(s, s.unclip_target_resource)
+        if item_name is None:
+            return False
+
+        has_item = getattr(s, item_name, 0) > 0
+        return has_item
+
+    def _get_unclip_info(
+        self, s: UnclippingAgentState, resource: Optional[str]
+    ) -> Optional[tuple[str, dict[str, int]]]:
+        """Get unclip item name and craft recipe for a resource type."""
+        if resource is None:
+            return None
+        return s.unclip_recipes.get(resource)
+
+    def _clear_unclip_state(self, s: UnclippingAgentState) -> None:
+        s.blocked_by_clipped_extractor = None
+        s.unclip_target_resource = None
+
+    def _set_unclip_state(self, s: UnclippingAgentState, resource_type: str, extractor: ExtractorInfo) -> None:
+        # Only set unclip state if we've discovered the recipe for this resource
+        if resource_type not in s.unclip_recipes:
+            return
+        s.blocked_by_clipped_extractor = extractor.position
+        s.unclip_target_resource = resource_type
+
+    def _get_vibe_for_phase(self, phase: Phase, state: UnclippingAgentState) -> str:
+        """Override to set correct vibe for CRAFT_UNCLIP and UNCLIP phases."""
+        # For crafting unclip items at the hub, use "gear" vibe
+        if phase == Phase.CRAFT_UNCLIP:
+            return "gear"
+
+        # For unclipping extractors, use "gear" vibe
+        if phase == Phase.UNCLIP:
+            return "gear"
+
+        # Otherwise use baseline logic
+        return super()._get_vibe_for_phase(phase, state)
+
+    def _update_phase(self, s: UnclippingAgentState) -> None:
+        """Override to add unclipping phase priorities before gathering."""
+        old_phase = s.phase
+
+        # Priority 1-2: Recharge and Deliver (handled by parent)
+        if s.energy < self._hyperparams.recharge_threshold_low or s.phase == Phase.RECHARGE or s.hearts > 0:
+            super()._update_phase(s)
+            if old_phase != s.phase:
+                s.cached_path = None
+                s.cached_path_target = None
+            return
+
+        # Priority 3: Assemble if possible (and not currently blocked by clipped resource)
+        heart_recipe = s.heart_recipe or {}
+        can_assemble = all(
+            getattr(s, res, 0) >= heart_recipe.get(res, 0) for res in ("carbon", "oxygen", "germanium", "silicon")
+        )
+
+        if can_assemble and s.blocked_by_clipped_extractor is None:
+            if s.phase != Phase.ASSEMBLE:
+                s.phase = Phase.ASSEMBLE
+                s.pending_use_resource = None
+                s.pending_use_amount = 0
+                s.waiting_at_extractor = None
+            if old_phase != s.phase:
+                s.cached_path = None
+                s.cached_path_target = None
+            return
+
+        # Priority 4: Unclipping workflow (before gathering)
+        if s.blocked_by_clipped_extractor is not None and s.unclip_target_resource is not None:
+            # First, check if the extractor is still clipped
+            target_pos = s.blocked_by_clipped_extractor
+            extractors = s.extractors.get(s.unclip_target_resource, [])
+            target_extractor = None
+            for ext in extractors:
+                if ext.position == target_pos:
+                    target_extractor = ext
+                    break
+
+            # If extractor is no longer clipped or not found, clear unclip state
+            if target_extractor is None or not target_extractor.clipped:
+                self._clear_unclip_state(s)
+                # Fall through to normal phase logic
+            else:
+                # Still clipped, continue unclipping workflow
+                info = self._get_unclip_info(s, s.unclip_target_resource)
+
+                if info is not None:
+                    item_name, craft_recipe = info
+                    item_count = getattr(s, item_name, 0)
+
+                    if item_count > 0:
+                        if s.phase != Phase.UNCLIP:
+                            s.phase = Phase.UNCLIP
+                        return
+                    elif craft_recipe and all(getattr(s, res, 0) >= amt for res, amt in craft_recipe.items()):
+                        if s.phase != Phase.CRAFT_UNCLIP:
+                            s.phase = Phase.CRAFT_UNCLIP
+                        return
+                    else:
+                        # Need to gather craft resources
+                        if s.phase != Phase.GATHER:
+                            s.phase = Phase.GATHER
+                        if old_phase != s.phase:
+                            s.cached_path = None
+                            s.cached_path_target = None
+                        return
+
+        # Priority 5: Default to GATHER (handled by parent)
+        super()._update_phase(s)
+        if old_phase != s.phase:
+            s.cached_path = None
+            s.cached_path_target = None
+
+    def _find_any_needed_extractor(self, s: UnclippingAgentState) -> Optional[tuple[ExtractorInfo, str]]:
+        """
+        Override to detect clipped extractors and trigger unclipping workflow.
+        """
+        # Try baseline logic first
+        result = super()._find_any_needed_extractor(s)
+        if result is not None:
+            self._clear_unclip_state(s)
+            return result
+
+        deficits = self._calculate_deficits(s)
+
+        def distance(pos: tuple[int, int]) -> int:
+            return abs(pos[0] - s.row) + abs(pos[1] - s.col)
+
+        # If we are already blocked, ensure we gather craft resources if needed
+        if s.blocked_by_clipped_extractor is not None:
+            info = self._get_unclip_info(s, s.unclip_target_resource)
+            if info is not None:
+                item_name, craft_recipe = info
+                if getattr(s, item_name, 0) > 0:
+                    return None  # Ready to unclip
+
+                # Need craft resources before crafting
+                if craft_recipe:
+                    for craft_resource, needed_amount in craft_recipe.items():
+                        if getattr(s, craft_resource, 0) < needed_amount:
+                            craft_extractors = s.extractors.get(craft_resource, [])
+                            available = [e for e in craft_extractors if not e.clipped and e.remaining_uses > 0]
+                            if available:
+                                nearest = min(available, key=lambda e: distance(e.position))
+                                return (nearest, craft_resource)
+
+                            clipped = [e for e in craft_extractors if e.clipped and e.remaining_uses > 0]
+                            if clipped:
+                                nearest = min(clipped, key=lambda e: distance(e.position))
+                                self._set_unclip_state(s, craft_resource, nearest)
+                            return None
+
+        # Look for deficits that are blocked by clipped extractors
+        for resource_type in self._priority_order:
+            if deficits.get(resource_type, 0) <= 0:
+                continue
+
+            extractors = s.extractors.get(resource_type, [])
+            if not extractors:
+                continue
+
+            available = [e for e in extractors if not e.clipped and e.remaining_uses > 0]
+            if available:
+                nearest = min(available, key=lambda e: distance(e.position))
+                return (nearest, resource_type)
+
+            clipped = [e for e in extractors if e.clipped and e.remaining_uses > 0]
+            if clipped:
+                nearest = min(clipped, key=lambda e: distance(e.position))
+                self._set_unclip_state(s, resource_type, nearest)
+                return None
+
+        # No special action
+        self._clear_unclip_state(s)
+        return None
+
+    def step_with_state(
+        self, obs: AgentObservation, state: UnclippingAgentState
+    ) -> tuple[Action, UnclippingAgentState]:
+        """Override to discover unclip recipes from observations."""
+        # First, let the base class handle observation parsing and state updates
+        action, updated_state = super().step_with_state(obs, state)
+        state = cast(UnclippingAgentState, updated_state)
+
+        # Now discover unclip recipes from the parsed observation
+        # We need to parse again to get the nearby_objects
+        parsed = self.parse_observation(state, obs)
+        self._discover_unclip_recipe(state, parsed)
+
+        return action, state
+
+    def _execute_phase(self, s: UnclippingAgentState) -> Action:
+        """Override to handle CRAFT_UNCLIP and UNCLIP phases."""
+        if s.phase == Phase.CRAFT_UNCLIP:
+            return self._do_craft_unclip(s)
+        elif s.phase == Phase.UNCLIP:
+            return self._do_unclip(s)
+        # All other phases handled by parent (GATHER, ASSEMBLE, DELIVER, RECHARGE)
+        return super()._execute_phase(s)
+
+    def _do_craft_unclip(self, s: UnclippingAgentState) -> Action:
+        """Craft unclip item at hub."""
+        info = self._get_unclip_info(s, s.unclip_target_resource)
+        if info is None:
+            self._clear_unclip_state(s)
+            s.phase = Phase.GATHER
+            return Action(name="noop")
+
+        item_name, craft_recipe = info
+
+        # Check if we have all craft resources
+        if craft_recipe and not all(getattr(s, res, 0) >= amt for res, amt in craft_recipe.items()):
+            # Need to gather craft resources first
+            s.phase = Phase.GATHER
+            return Action(name="noop")
+
+        # Explore until we find hub
+        explore_action = self._explore_until(
+            s, condition=lambda: s.stations.get("hub") is not None, reason="Need hub for crafting"
+        )
+        if explore_action is not None:
+            return explore_action
+
+        # Vibe is automatically set by _get_vibe_for_phase to the input resource (e.g., "carbon" for decoder)
+
+        hub = s.stations.get("hub")
+        if hub is None:
+            return Action(name="noop")
+
+        ar, ac = hub
+        if is_adjacent((s.row, s.col), hub):
+            return use_object_at(s, hub)
+
+        return self._move_towards(s, hub, reach_adjacent=True)
+
+    def _do_unclip(self, s: UnclippingAgentState) -> Action:
+        """Use unclip item on clipped extractor."""
+        if s.blocked_by_clipped_extractor is None:
+            s.phase = Phase.GATHER
+            self._clear_unclip_state(s)
+            return Action(name="noop")
+
+        info = self._get_unclip_info(s, s.unclip_target_resource)
+        if info is None:
+            self._clear_unclip_state(s)
+            s.phase = Phase.GATHER
+            return Action(name="noop")
+
+        item_name, _ = info
+        if getattr(s, item_name, 0) <= 0:
+            # Lost the item before reaching extractor
+            self._clear_unclip_state(s)
+            s.phase = Phase.GATHER
+            return Action(name="noop")
+
+        # Navigate to clipped extractor
+        target = s.blocked_by_clipped_extractor
+        tr, tc = target
+        is_at_target = s.row == tr and s.col == tc
+
+        if is_at_target:
+            # Already on the extractor - it should be unclipped now
+            # Wait for next step to verify and clear state
+            return Action(name="noop")
+
+        if is_adjacent((s.row, s.col), target):
+            # Adjacent to clipped extractor - use it to unclip (like using any other object)
+            action = use_object_at(s, target)
+            # Don't clear unclip state yet - wait until next step to verify it worked
+            # The state will be cleared in _update_phase when we see the extractor is unclipped
+            return action
+
+        # Not adjacent yet, move towards it
+        return self._move_towards(s, target, reach_adjacent=True)
+
+
+# ============================================================================
+# Policy Wrapper Class
+# ============================================================================
+
+
+class UnclippingPolicy(MultiAgentPolicy):
+    """Multi-agent policy wrapper for UnclippingAgent.
+
+    This class wraps UnclippingAgent to work with the policy interface.
+    It handles multiple agents, each with their own UnclippingAgent instance.
+    """
+
+    short_names = ["ladybug_py"]
+
+    def __init__(
+        self,
+        policy_env_info: PolicyEnvInterface,
+        device: str = "cpu",
+        hyperparams: Optional[UnclippingHyperparameters] = None,
+    ):
+        super().__init__(policy_env_info, device=device)
+        self._agent_policies: dict[int, StatefulAgentPolicy[UnclippingAgentState]] = {}
+        self._hyperparams = hyperparams or UnclippingHyperparameters()
+
+    def agent_policy(self, agent_id: int) -> StatefulAgentPolicy[UnclippingAgentState]:
+        if agent_id not in self._agent_policies:
+            # UnclippingAgentPolicyImpl uses UnclippingAgentState but inherits from
+            # BaselineAgentPolicyImpl typed with SimpleAgentState, requiring a cast
+            policy = cast(
+                StatefulAgentPolicy[UnclippingAgentState],
+                StatefulAgentPolicy(
+                    UnclippingAgentPolicyImpl(self._policy_env_info, agent_id, self._hyperparams),
+                    self._policy_env_info,
+                    agent_id=agent_id,
+                ),
+            )
+            self._agent_policies[agent_id] = policy
+        return self._agent_policies[agent_id]