cogames 0.3.49__py3-none-any.whl → 0.3.64__py3-none-any.whl

cogames/policy/scripted_agent/demo_policy.py

@@ -1,244 +0,0 @@
-import random
-
-from mettagrid.policy.policy import MultiAgentPolicy, StatefulAgentPolicy, StatefulPolicyImpl
-from mettagrid.simulator.interface import AgentObservation
-
-from .types import BaselineHyperparameters, CellType, SimpleAgentState
-from .utils import (
-    change_vibe_action,
-    is_station,
-    is_wall,
-    parse_observation,
-    read_inventory_from_obs,
-    update_agent_position,
-    use_object_at,
-)
-
-
-def manhattan(a, b):
-    return abs(a[0] - b[0]) + abs(a[1] - b[1])
-
-
-class DemoPolicyImpl(StatefulPolicyImpl[SimpleAgentState]):
-    def __init__(self, policy_env_info, agent_id, hyperparams, *, heart_recipe=None):
-        self._agent_id = agent_id
-        self._hyperparams = hyperparams
-        self._policy_env_info = policy_env_info
-        self._actions = policy_env_info.actions
-        self._move_deltas = {"north": (-1, 0), "south": (1, 0), "east": (0, 1), "west": (0, -1)}
-
-        self._obs_hr = policy_env_info.obs_height // 2
-        self._obs_wr = policy_env_info.obs_width // 2
-
-        if heart_recipe:
-            self._heart_recipe = {k: v for k, v in heart_recipe.items() if k != "energy"}
-        else:
-            self._heart_recipe = None
-
-    def initial_agent_state(self):
-        center = 25
-        return SimpleAgentState(
-            agent_id=self._agent_id,
-            map_height=50,
-            map_width=50,
-            occupancy=[[CellType.FREE.value] * 50 for _ in range(50)],
-            row=center,
-            col=center,
-            heart_recipe=self._heart_recipe,
-        )
-
-    # ------------------------------------------------------------
-    # Utility helpers (kept tiny)
-    # ------------------------------------------------------------
-
-    def _adjacent(self, s, pos):
-        return manhattan((s.row, s.col), pos) == 1
-
-    def _random_step(self, s, parsed):
-        dirs = list(self._move_deltas.keys())
-        random.shuffle(dirs)
-        blocked = {
-            (r, c)
-            for (r, c), obj in parsed.nearby_objects.items()
-            if self._adjacent(s, (r, c))
-            and (
-                is_wall(obj.name)
-                or "extractor" in obj.name
-                or is_station(obj.name, "assembler")
-                or is_station(obj.name, "chest")
-                or is_station(obj.name, "charger")
-                or (obj.name == "agent" and obj.agent_group != s.agent_id)
-            )
-        }
-        for d in dirs:
-            dr, dc = self._move_deltas[d]
-            nr, nc = s.row + dr, s.col + dc
-            if (nr, nc) not in blocked:
-                return self._actions.move.Move(d)
-        return self._actions.noop.Noop()
-
-    def _step_towards(self, s, target, parsed):
-        """Single-step greedy pursuit, else random."""
-        r, c = s.row, s.col
-        tr, tc = target
-        cand = []
-        if abs(tr - r) >= abs(tc - c):
-            if tr < r:
-                cand.append("north")
-            elif tr > r:
-                cand.append("south")
-            if tc < c:
-                cand.append("west")
-            elif tc > c:
-                cand.append("east")
-        else:
-            if tc < c:
-                cand.append("west")
-            elif tc > c:
-                cand.append("east")
-            if tr < r:
-                cand.append("north")
-            elif tr > r:
-                cand.append("south")
-
-        blocked = {
-            (rr, cc)
-            for (rr, cc), obj in parsed.nearby_objects.items()
-            if self._adjacent(s, (rr, cc))
-            and (
-                is_wall(obj.name)
-                or "extractor" in obj.name
-                or is_station(obj.name, "assembler")
-                or is_station(obj.name, "chest")
-                or is_station(obj.name, "charger")
-                or (obj.name == "agent" and obj.agent_group != s.agent_id)
-            )
-        }
-
-        for d in cand:
-            dr, dc = self._move_deltas[d]
-            nr, nc = r + dr, c + dc
-            if (nr, nc) not in blocked:
-                return self._actions.move.Move(d)
-
-        return self._random_step(s, parsed)
-
-    def _closest(self, s, parsed, pred):
-        items = [pos for pos, obj in parsed.nearby_objects.items() if pred(obj)]
-        return min(items, key=lambda p: manhattan((s.row, s.col), p)) if items else None
-
-    def _rtype(self, name):
-        name = name.lower().replace("clipped_", "")
-        if "_extractor" not in name:
-            return None
-        name = name.replace("_extractor", "")
-        return name if name in ("carbon", "oxygen", "germanium", "silicon") else None
-
-    # ------------------------------------------------------------
-    # Main step
-    # ------------------------------------------------------------
-    def step_with_state(self, obs: AgentObservation, s: SimpleAgentState):
-        s.step_count += 1
-        read_inventory_from_obs(s, obs, obs_hr=self._obs_hr, obs_wr=self._obs_wr)
-        update_agent_position(s, move_deltas=self._move_deltas)
-
-        parsed = parse_observation(
-            s,
-            obs,
-            obs_hr=self._obs_hr,
-            obs_wr=self._obs_wr,
-            spatial_feature_names={"tag", "cooldown_remaining", "clipped", "remaining_uses"},
-            agent_feature_key_by_name={"agent:group": "agent_group", "agent:frozen": "agent_frozen"},
-            protocol_input_prefix="protocol_input:",
-            protocol_output_prefix="protocol_output:",
-            tag_names=self._policy_env_info.tag_id_to_name,
-        )
-
-        # Learn recipe if visible
-        if s.heart_recipe is None:
-            for _pos, obj in parsed.nearby_objects.items():
-                if obj.name == "assembler" and obj.protocol_outputs.get("heart", 0) > 0:
-                    s.heart_recipe = {k: v for k, v in obj.protocol_inputs.items() if k != "energy"}
-
-        # ---------------- PRE-PHASE: find assembler to learn recipe ----------------
-        if s.heart_recipe is None:
-            if s.current_glyph != "heart_a":
-                s.current_glyph = "heart_a"
-                return change_vibe_action("heart_a", actions=self._actions), s
-
-            assembler = self._closest(s, parsed, lambda o: is_station(o.name.lower(), "assembler"))
-            if assembler:
-                if self._adjacent(s, assembler):
-                    return use_object_at(s, assembler, actions=self._actions, move_deltas=self._move_deltas), s
-                return self._step_towards(s, assembler, parsed), s
-
-            return self._random_step(s, parsed), s
-
-        # ---------------- MAIN PHASE ----------------
-
-        # Deliver hearts
-        if s.hearts > 0:
-            chest = self._closest(s, parsed, lambda o: is_station(o.name.lower(), "chest"))
-            if chest:
-                if s.current_glyph != "default":
-                    s.current_glyph = "default"
-                    return change_vibe_action("default", actions=self._actions), s
-                if self._adjacent(s, chest):
-                    return use_object_at(s, chest, actions=self._actions, move_deltas=self._move_deltas), s
-                return self._step_towards(s, chest, parsed), s
-
-        # Assemble
-        if (
-            s.carbon >= s.heart_recipe.get("carbon", 0)
-            and s.oxygen >= s.heart_recipe.get("oxygen", 0)
-            and s.germanium >= s.heart_recipe.get("germanium", 0)
-            and s.silicon >= s.heart_recipe.get("silicon", 0)
-        ):
-            assembler = self._closest(s, parsed, lambda o: is_station(o.name.lower(), "assembler"))
-            if assembler:
-                if s.current_glyph != "heart_a":
-                    s.current_glyph = "heart_a"
-                    return change_vibe_action("heart_a", actions=self._actions), s
-                if self._adjacent(s, assembler):
-                    return use_object_at(s, assembler, actions=self._actions, move_deltas=self._move_deltas), s
-                return self._step_towards(s, assembler, parsed), s
-
-        # Gather needed resources
-        deficits = {
-            r: s.heart_recipe.get(r, 0) - getattr(s, r, 0) for r in ("carbon", "oxygen", "germanium", "silicon")
-        }
-        needed = [
-            (pos, obj, self._rtype(obj.name.lower()))
-            for pos, obj in parsed.nearby_objects.items()
-            if "extractor" in obj.name.lower()
-        ]
-
-        needed = [(pos, obj, r) for pos, obj, r in needed if r and deficits[r] > 0]
-
-        if needed:
-            pos, obj, r = min(needed, key=lambda x: manhattan((s.row, s.col), x[0]))
-            if self._adjacent(s, pos):
-                return use_object_at(s, pos, actions=self._actions, move_deltas=self._move_deltas), s
-            return self._step_towards(s, pos, parsed), s
-
-        # Otherwise wander
-        return self._random_step(s, parsed), s
-
-
-class DemoPolicy(MultiAgentPolicy):
-    short_names = ["tiny_baseline"]
-
-    def __init__(self, policy_env_info, device: str = "cpu", hyperparams=None, *, heart_recipe=None):
-        super().__init__(policy_env_info, device=device)
-        self._hyperparams = hyperparams or BaselineHyperparameters()
-        self._heart_recipe = heart_recipe
-        self._agent_policies = {}
-
-    def agent_policy(self, agent_id):
-        if agent_id not in self._agent_policies:
-            self._agent_policies[agent_id] = StatefulAgentPolicy(
-                DemoPolicyImpl(self._policy_env_info, agent_id, self._hyperparams, heart_recipe=self._heart_recipe),
-                self._policy_env_info,
-                agent_id=agent_id,
-            )
-        return self._agent_policies[agent_id]

cogames/policy/scripted_agent/pathfinding.py

@@ -1,126 +0,0 @@
-"""
-Pathfinding utilities for scripted agents.
-
-This module contains A* pathfinding implementation and related utilities
-for navigating the grid world.
-"""
-
-from __future__ import annotations
-
-from collections import deque
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from cogames.policy.scripted_agent.types import CellType, SimpleAgentState
-
-
-def compute_goal_cells(
-    state: SimpleAgentState, target: tuple[int, int], reach_adjacent: bool, cell_type: type[CellType]
-) -> list[tuple[int, int]]:
-    """
-    Compute the set of goal cells for pathfinding.
-    """
-    if not reach_adjacent:
-        return [target]
-
-    goals: list[tuple[int, int]] = []
-    for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
-        nr, nc = target[0] + dr, target[1] + dc
-        if is_traversable(state, nr, nc, cell_type):
-            goals.append((nr, nc))
-
-    # If no adjacent traversable tiles are known yet, allow exploring toward unknown ones
-    if not goals:
-        for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
-            nr, nc = target[0] + dr, target[1] + dc
-            if is_within_bounds(state, nr, nc) and state.occupancy[nr][nc] != cell_type.OBSTACLE.value:
-                goals.append((nr, nc))
-    return goals
-
-
-def shortest_path(
-    state: SimpleAgentState,
-    start: tuple[int, int],
-    goals: list[tuple[int, int]],
-    allow_goal_block: bool,
-    cell_type: type[CellType],
-) -> list[tuple[int, int]]:
-    """
-    Find shortest path from start to any goal using BFS.
-    """
-    goal_set = set(goals)
-    queue: deque[tuple[int, int]] = deque([start])
-    came_from: dict[tuple[int, int], tuple[int, int] | None] = {start: None}
-
-    def walkable(r: int, c: int) -> bool:
-        if (r, c) in goal_set and allow_goal_block:
-            return True
-        return is_traversable(state, r, c, cell_type)
-
-    while queue:
-        current = queue.popleft()
-        if current in goal_set:
-            return reconstruct_path(came_from, current)
-
-        for nr, nc in get_neighbors(state, current):
-            if (nr, nc) not in came_from and walkable(nr, nc):
-                came_from[(nr, nc)] = current
-                queue.append((nr, nc))
-
-    return []
-
-
-def reconstruct_path(
-    came_from: dict[tuple[int, int], tuple[int, int] | None],
-    current: tuple[int, int],
-) -> list[tuple[int, int]]:
-    """
-    Reconstruct path from BFS came_from dict.
-    """
-    path: list[tuple[int, int]] = []
-    while came_from[current] is not None:
-        path.append(current)
-        prev = came_from[current]
-        assert prev is not None  # Loop condition ensures this
-        current = prev
-    path.reverse()
-    return path
-
-
-def get_neighbors(state: SimpleAgentState, pos: tuple[int, int]) -> list[tuple[int, int]]:
-    """
-    Get valid neighboring positions (4-connected grid).
-    """
-    r, c = pos
-    candidates = [(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)]
-    return [(nr, nc) for nr, nc in candidates if is_within_bounds(state, nr, nc)]
-
-
-def is_within_bounds(state: SimpleAgentState, r: int, c: int) -> bool:
-    """
-    Check if position is within map bounds.
-    """
-    return 0 <= r < state.map_height and 0 <= c < state.map_width
-
-
-def is_passable(state: SimpleAgentState, r: int, c: int, cell_type: type[CellType]) -> bool:
-    """
-    Check if a cell is passable (not an obstacle).
-    """
-    if not is_within_bounds(state, r, c):
-        return False
-    return is_traversable(state, r, c, cell_type)
-
-
-def is_traversable(state: SimpleAgentState, r: int, c: int, cell_type: type[CellType]) -> bool:
-    """
-    Check if a cell is traversable (free and no agent there).
-    """
-    if not is_within_bounds(state, r, c):
-        return False
-    # Don't walk through other agents
-    if (r, c) in state.agent_occupancy:
-        return False
-    cell = state.occupancy[r][c]
-    # Only traverse cells we KNOW are free, not unknown cells
-    return cell == cell_type.FREE.value

cogames/policy/scripted_agent/starter_agent.py

@@ -1,136 +0,0 @@
-"""
-Sample Policy for the Cogames environment.
-
-Helps a Cog move carbon from an extractor to a chest.
-
-Note to users of this policy:
-We don't intend for scripted policies to be the final word on how policies are generated (e.g., we expect the
-environment to be complicated enough that trained agents will be necessary). So we expect that scripting policies
-is a good way to start, but don't want you to get stuck here. Feel free to prove us wrong!
-
-Note to cogames developers:
-This policy should be kept relatively minimalist, without dependencies on intricate algorithms.
-"""
-
-from __future__ import annotations
-
-import random
-from dataclasses import dataclass
-from typing import Tuple
-
-from mettagrid.policy.policy import MultiAgentPolicy, StatefulAgentPolicy, StatefulPolicyImpl
-from mettagrid.policy.policy_env_interface import PolicyEnvInterface
-from mettagrid.simulator import Action
-from mettagrid.simulator.interface import AgentObservation
-
-
-@dataclass
-class StarterCogState:
-    target_location: tuple[int, int] | None = None
-    chest_location: tuple[int, int] | None = None
-    resource_to_collect: str = "carbon"
-    # Location of an extractor for the target resource
-    extractor_location: tuple[int, int] | None = None
-    # Current position relative to the starting position.
-    # We expect some moves to fail, so all positions should be treated somewhat loosely.
-    position: tuple[int, int] = (0, 0)
-    have_inventory: bool = False
-
-
-class StarterCogPolicyImpl(StatefulPolicyImpl[StarterCogState]):
-    def __init__(
-        self,
-        policy_env_info: PolicyEnvInterface,
-        agent_id: int,
-    ):
-        self._agent_id = agent_id
-        self._policy_env_info = policy_env_info
-
-        # Action lookup
-        self._actions = policy_env_info.actions
-
-    def _parse_observation(self, obs: AgentObservation, state: StarterCogState) -> StarterCogState:
-        """Parse the observation and update the state."""
-
-        extractor_tag_value = self._policy_env_info.tags.index(f"{state.resource_to_collect}_extractor")
-        chest_tag_value = self._policy_env_info.tags.index("chest")
-        state.have_inventory = False
-        for token in obs.tokens:
-            if token.feature.name == "last_action":
-                # Update our current (relative) position.
-                # TODO: This is wrong if we moved to interact with a station.
-                if token.value == self._policy_env_info.action_names.index("move_north"):
-                    state.position = (state.position[0] - 1, state.position[1])
-                elif token.value == self._policy_env_info.action_names.index("move_south"):
-                    state.position = (state.position[0] + 1, state.position[1])
-                elif token.value == self._policy_env_info.action_names.index("move_west"):
-                    state.position = (state.position[0], state.position[1] - 1)
-                elif token.value == self._policy_env_info.action_names.index("move_east"):
-                    state.position = (state.position[0], state.position[1] + 1)
-                break
-        for token in obs.tokens:
-            if token.location == (5, 5):
-                if token.feature.name == f"inv:{state.resource_to_collect}":
-                    state.have_inventory = True
-                continue
-            token_location = (state.position[0] + token.location[0] - 5, state.position[1] + token.location[1] - 5)
-            if token.feature.name == "tag":
-                if token.value == extractor_tag_value:
-                    state.extractor_location = token_location
-                elif token.value == chest_tag_value:
-                    state.chest_location = token_location
-            # It would probably be a good idea to keep track of obstacles
-        return state
-
-    def _go_to(self, state: StarterCogState, target: tuple[int, int]) -> Tuple[Action, StarterCogState]:
-        """Go to the target location."""
-        # Let's just go straight there! I hope we don't run into anything and get stuck.
-        possible_actions = []
-        if state.position[0] < target[0]:
-            possible_actions.append("move_south")
-        if state.position[0] > target[0]:
-            possible_actions.append("move_north")
-        if state.position[1] < target[1]:
-            possible_actions.append("move_east")
-        if state.position[1] > target[1]:
-            possible_actions.append("move_west")
-        action = Action(name=random.choice(possible_actions))
-        return action, state
-
-    def step_with_state(self, obs: AgentObservation, state: StarterCogState) -> Tuple[Action, StarterCogState]:
-        """Compute the action for this Cog."""
-        state = self._parse_observation(obs, state)
-        if state.have_inventory:
-            if state.chest_location is not None:
-                return self._go_to(state, state.chest_location)
-        else:
-            if state.extractor_location is not None:
-                return self._go_to(state, state.extractor_location)
-        direction = random.choice(["north", "south", "east", "west"])
-        return Action(name="move_" + direction), state
-
-    def initial_agent_state(self) -> StarterCogState:
-        """Get the initial state for a new agent."""
-        return StarterCogState()
-
-
-# ============================================================================
-# Policy Wrapper Classes
-# ============================================================================
-
-
-class StarterPolicy(MultiAgentPolicy):
-    # short_names = ["scripted_starter"]  # Uncomment to register a shorthand
-
-    def __init__(self, policy_env_info: PolicyEnvInterface, device: str = "cpu"):
-        super().__init__(policy_env_info, device=device)
-        self._agent_policies: dict[int, StatefulAgentPolicy[StarterCogState]] = {}
-
-    def agent_policy(self, agent_id: int) -> StatefulAgentPolicy[StarterCogState]:
-        if agent_id not in self._agent_policies:
-            self._agent_policies[agent_id] = StatefulAgentPolicy(
-                StarterCogPolicyImpl(self._policy_env_info, agent_id),
-                self._policy_env_info,
-                agent_id=agent_id,
-            )
-        return self._agent_policies[agent_id]