cogames 0.3.59.post1.dev2__py3-none-any.whl → 0.3.65__py3-none-any.whl

cogames/cogs_vs_clips/evals/spanning_evals.py

@@ -5,26 +5,24 @@ from __future__ import annotations
 
 import logging
 
-from cogames.cogs_vs_clips.mission import Mission, Site
-from cogames.cogs_vs_clips.procedural import MachinaArena
+from cogames.cogs_vs_clips.mission import CvCMission
 from cogames.cogs_vs_clips.sites import HELLO_WORLD, TRAINING_FACILITY
+from cogames.cogs_vs_clips.terrain import MachinaArena
 from cogames.cogs_vs_clips.variants import (
-    CompassVariant,
     DarkSideVariant,
     DistantResourcesVariant,
     EmptyBaseVariant,
     EnergizedVariant,
-    PackRatVariant,
     QuadrantBuildingsVariant,
-    RoughTerrainVariant,
     SingleResourceUniformVariant,
     SuperChargedVariant,
 )
+from cogames.core import CoGameSite
 from mettagrid.mapgen.mapgen import MapGen
 
 logger = logging.getLogger(__name__)
 
-SMALL_HELLO_WORLD = Site(
+SMALL_HELLO_WORLD = CoGameSite(
     name="small_hello_world",
     description="Small hello world map.",
     map_builder=MapGen.Config(width=50, height=50, instance=MachinaArena.Config(spawn_count=20)),
@@ -32,7 +30,7 @@ SMALL_HELLO_WORLD = Site(
     max_cogs=20,
 )
 
-MEDIUM_HELLO_WORLD = Site(
+MEDIUM_HELLO_WORLD = CoGameSite(
     name="medium_hello_world",
     description="Medium hello world map.",
     map_builder=MapGen.Config(width=100, height=100, instance=MachinaArena.Config(spawn_count=20)),
@@ -40,7 +38,7 @@ MEDIUM_HELLO_WORLD = Site(
     max_cogs=20,
 )
 
-LARGE_HELLO_WORLD = Site(
+LARGE_HELLO_WORLD = CoGameSite(
     name="large_hello_world",
     description="Large hello world map.",
     map_builder=MapGen.Config(width=500, height=500, instance=MachinaArena.Config(spawn_count=20)),
@@ -49,19 +47,19 @@ LARGE_HELLO_WORLD = Site(
 )
 
 # Resource Bottleneck evals
-OxygenBottleneck = Mission(
+OxygenBottleneck = CvCMission(
     name="oxygen_bottleneck",
     description="Oxygen is the limiting resource; agents must prioritize oxygen over other resources.",
     site=HELLO_WORLD,
     variants=[
         EmptyBaseVariant(missing=["oxygen_extractor"]),
         SingleResourceUniformVariant(building_name="oxygen_extractor"),
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
 # Energy Starved evals
-EnergyStarved = Mission(
+EnergyStarved = CvCMission(
     name="energy_starved",
     description="Energy is the limiting resource; agents must prioritize energy over other resources.",
     site=HELLO_WORLD,
@@ -74,17 +72,17 @@ EnergyStarved = Mission(
 # Curated difficulty tiers per mission
 # ------------------------------------------------------------
 # Oxygen Bottleneck
-OxygenBottleneckEasy = Mission(
+OxygenBottleneckEasy = CvCMission(
     name="oxygen_bottleneck_easy",
     description="Easy: tuned oxygen focus with simple layout and generous capacities.",
     site=HELLO_WORLD,
     variants=[
         SingleResourceUniformVariant(building_name="oxygen_extractor"),
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
-OxygenBottleneckStandard = Mission(
+OxygenBottleneckStandard = CvCMission(
     name="oxygen_bottleneck_standard",
     description="Standard: oxygen is the bottleneck; extractor missing at base.",
     site=HELLO_WORLD,
@@ -93,18 +91,18 @@ OxygenBottleneckStandard = Mission(
     ],
 )
 
-OxygenBottleneckHard = Mission(
+OxygenBottleneckHard = CvCMission(
     name="oxygen_bottleneck_hard",
-    description="Hard: oxygen bottleneck plus rough terrain.",
+    description="Hard: oxygen bottleneck with dark side.",
     site=HELLO_WORLD,
     variants=[
         EmptyBaseVariant(missing=["oxygen_extractor"]),
-        RoughTerrainVariant(),
+        DarkSideVariant(),
     ],
 )
 
 # Energy Starved
-EnergyStarvedEasy = Mission(
+EnergyStarvedEasy = CvCMission(
     name="energy_starved_easy",
     description="Easy: abundant energy regen and capacity.",
     site=HELLO_WORLD,
@@ -114,7 +112,7 @@ EnergyStarvedEasy = Mission(
     ],
 )
 
-EnergyStarvedStandard = Mission(
+EnergyStarvedStandard = CvCMission(
     name="energy_starved_standard",
     description="Standard: energy is the limiting resource with dark-side regen.",
     site=HELLO_WORLD,
@@ -123,18 +121,17 @@ EnergyStarvedStandard = Mission(
     ],
 )
 
-EnergyStarvedHard = Mission(
+EnergyStarvedHard = CvCMission(
     name="energy_starved_hard",
-    description="Hard: energy bottleneck with dark side and rough terrain.",
+    description="Hard: energy bottleneck with dark side.",
     site=HELLO_WORLD,
     variants=[
         DarkSideVariant(),
-        RoughTerrainVariant(),
     ],
 )
 
 # Collect Distant Resources evals
-DistantResources = Mission(
+DistantResources = CvCMission(
     name="distant_resources",
     description="Resources scattered far from base; heavy routing coordination.",
     site=HELLO_WORLD,
@@ -145,39 +142,37 @@ DistantResources = Mission(
 )
 
 # Distant Resources tiers
-DistantResourcesEasy = Mission(
+DistantResourcesEasy = CvCMission(
     name="distant_resources_easy",
-    description="Easy: simplified distribution and navigation aids.",
+    description="Easy: simplified distribution with generous capacity.",
     site=HELLO_WORLD,
     variants=[
-        CompassVariant(),
-        PackRatVariant(),
+        EnergizedVariant(),
         DistantResourcesVariant(),
     ],
 )
 
-DistantResourcesStandard = Mission(
+DistantResourcesStandard = CvCMission(
     name="distant_resources_standard",
     description="Standard: resources scattered far from base.",
     site=HELLO_WORLD,
     variants=[
-        CompassVariant(),
         DistantResourcesVariant(),
     ],
 )
 
-DistantResourcesHard = Mission(
+DistantResourcesHard = CvCMission(
     name="distant_resources_hard",
-    description="Hard: distant resources with rough terrain.",
+    description="Hard: distant resources with dark side.",
     site=HELLO_WORLD,
     variants=[
         DistantResourcesVariant(),
-        RoughTerrainVariant(),
+        DarkSideVariant(),
     ],
 )
 
 # Divide and Conquer evals
-QuadrantBuildings = Mission(
+QuadrantBuildings = CvCMission(
     name="quadrant_buildings",
     description="Place buildings in the four quadrants of the map.",
     site=HELLO_WORLD,
@@ -188,18 +183,17 @@ QuadrantBuildings = Mission(
 )
 
 # Quadrant Buildings tiers
-QuadrantBuildingsEasy = Mission(
+QuadrantBuildingsEasy = CvCMission(
     name="quadrant_buildings_easy",
-    description="Easy: buildings in quadrants with navigation aid and inventory boost.",
+    description="Easy: buildings in quadrants with energy boost.",
     site=HELLO_WORLD,
     variants=[
         QuadrantBuildingsVariant(),
-        CompassVariant(),
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
-QuadrantBuildingsStandard = Mission(
+QuadrantBuildingsStandard = CvCMission(
     name="quadrant_buildings_standard",
     description="Standard: buildings placed in quadrants.",
     site=HELLO_WORLD,
@@ -209,54 +203,54 @@ QuadrantBuildingsStandard = Mission(
     ],
 )
 
-QuadrantBuildingsHard = Mission(
+QuadrantBuildingsHard = CvCMission(
     name="quadrant_buildings_hard",
-    description="Hard: quadrant distribution with empty base and rough terrain.",
+    description="Hard: quadrant distribution with empty base and dark side.",
     site=HELLO_WORLD,
     variants=[
         QuadrantBuildingsVariant(),
         EmptyBaseVariant(),
-        RoughTerrainVariant(),
+        DarkSideVariant(),
     ],
 )
 
-EasyHeartsTraining = Mission(
+EasyHeartsTraining = CvCMission(
     name="easy_hearts_training",
-    description="Simplified heart crafting with generous caps and extractor base.",
+    description="Simplified heart crafting with generous energy.",
     site=TRAINING_FACILITY,
     variants=[
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
-EasyHeartsSmallWorld = Mission(
+EasyHeartsSmallWorld = CvCMission(
     name="easy_small_hearts",
-    description="Simplified heart crafting with generous caps and extractor base.",
+    description="Simplified heart crafting with generous energy.",
     site=SMALL_HELLO_WORLD,
     variants=[
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
-EasyHeartsMediumWorld = Mission(
+EasyHeartsMediumWorld = CvCMission(
     name="easy_medium_hearts",
-    description="Simplified heart crafting with generous caps and extractor base.",
+    description="Simplified heart crafting with generous energy.",
     site=MEDIUM_HELLO_WORLD,
     variants=[
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
-EasyHeartsLargeWorld = Mission(
+EasyHeartsLargeWorld = CvCMission(
     name="easy_large_hearts",
-    description="Simplified heart crafting with generous caps and extractor base.",
+    description="Simplified heart crafting with generous energy.",
     site=LARGE_HELLO_WORLD,
     variants=[
-        PackRatVariant(),
+        EnergizedVariant(),
     ],
 )
 
-EVAL_MISSIONS: list[Mission] = [
+EVAL_MISSIONS: list[CvCMission] = [
     # Oxygen bottleneck tiers
     OxygenBottleneckEasy,
     OxygenBottleneckStandard,

cogames/cogs_vs_clips/mission.py

@@ -1,188 +1,70 @@
 from __future__ import annotations
 
-from abc import ABC
-from typing import TypeVar, override
-
 from pydantic import Field
-from typing_extensions import Self
 
+from cogames.cogs_vs_clips.clips import ClipsConfig
 from cogames.cogs_vs_clips.cog import CogConfig
+from cogames.cogs_vs_clips.config import CvCConfig
 from cogames.cogs_vs_clips.stations import (
-    ELEMENTS,
-    GEAR,
-    CogsGuardChestConfig,
+    CvCChestConfig,
+    CvCExtractorConfig,
+    CvCGearStationConfig,
+    CvCHubConfig,
+    CvCJunctionConfig,
     CvCWallConfig,
-    GearStationConfig,
-    HubConfig,
-    JunctionConfig,
-    SimpleExtractorConfig,
 )
-from mettagrid.base_config import Config
+from cogames.cogs_vs_clips.team import CogTeam
+from cogames.cogs_vs_clips.variants import NumCogsVariant
+from cogames.core import (
+    MAP_MISSION_DELIMITER,
+    CoGameMission,
+    CoGameMissionVariant,
+    CoGameSite,
+)
 from mettagrid.config.action_config import (
     ActionsConfig,
     ChangeVibeActionConfig,
     MoveActionConfig,
     NoopActionConfig,
 )
-from mettagrid.config.event_config import EventConfig, periodic
-from mettagrid.config.filter import isAlignedTo, isNear
 from mettagrid.config.game_value import inv
-from mettagrid.config.game_value import stat as game_stat
-from mettagrid.config.mettagrid_config import (
-    CollectiveConfig,
-    GameConfig,
-    InventoryConfig,
-    MettaGridConfig,
-    ResourceLimitsConfig,
-)
-from mettagrid.config.mutation import alignTo
+from mettagrid.config.mettagrid_config import GameConfig, MettaGridConfig
 from mettagrid.config.obs_config import GlobalObsConfig, ObsConfig
-from mettagrid.config.reward_config import numObjects, reward
-from mettagrid.config.tag import typeTag
-from mettagrid.config.vibes import Vibe
 from mettagrid.map_builder.map_builder import AnyMapBuilderConfig
 
-# Type variable for mission types
-TMission = TypeVar("TMission", bound="MissionBase")
-
-
-class MissionVariant(Config, ABC):
-    # Note: we could derive the name from the class name automatically, but it would make it
-    # harder to find the variant source code based on CLI interactions.
-    name: str
-    description: str = Field(default="")
-
-    def modify_mission(self, mission: MissionBase) -> None:
-        # Override this method to modify the mission.
-        # Variants are allowed to modify the mission in-place - it's guaranteed to be a one-time only instance.
-        pass
-
-    def modify_env(self, mission: MissionBase, env: MettaGridConfig) -> None:
-        # Override this method to modify the produced environment.
-        # Variants are allowed to modify the environment in-place.
-        pass
-
-    def compat(self, mission: MissionBase) -> bool:
-        """Check if this variant is compatible with the given mission.
-
-        Returns True if the variant can be safely applied to the mission.
-        Override this method to add compatibility checks.
-        """
-        return True
-
-    def apply(self, mission: TMission) -> TMission:
-        mission = mission.model_copy(deep=True)
-        mission.variants.append(self)
-        self.modify_mission(mission)
-        return mission
-
-    # Temporary helper useful as long as we have one-time variants in missions.py file.
-    def as_mission(self, name: str, description: str, site: Site) -> Mission:
-        return Mission(
-            name=name,
-            description=description,
-            site=site,
-            variants=[self],
-        )
-
-
-class NumCogsVariant(MissionVariant):
-    name: str = "num_cogs"
-    description: str = "Set the number of cogs for the mission."
-    num_cogs: int
-
-    @override
-    def modify_mission(self, mission: Mission) -> None:
-        if self.num_cogs < mission.site.min_cogs or self.num_cogs > mission.site.max_cogs:
-            raise ValueError(
-                f"Invalid number of cogs for {mission.site.name}: {self.num_cogs}. "
-                + f"Must be between {mission.site.min_cogs} and {mission.site.max_cogs}"
-            )
-
-        mission.num_cogs = self.num_cogs
-
-
-class Site(Config):
-    name: str
-    description: str
-    map_builder: AnyMapBuilderConfig
-
-    min_cogs: int = Field(default=1, ge=1)
-    max_cogs: int = Field(default=1000, ge=1)
-
-
-MAP_MISSION_DELIMITER = "."
-
-
-class MissionBase(Config, ABC):
-    """Base class for Mission configurations with common fields and methods."""
-
-    name: str
-    description: str
-    site: Site
-    num_cogs: int | None = None
-
-    # Variants are applied to the mission immediately, and to its env when make_env is called
-    variants: list[MissionVariant] = Field(default_factory=list)
-
-    max_steps: int = Field(default=10000)
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        # Can't call `variant.apply` here because it will create a new mission instance
-        for variant in self.variants:
-            variant.modify_mission(self)
-
-    def with_variants(self, variants: list[MissionVariant]) -> Self:
-        mission = self
-        for variant in variants:
-            mission = variant.apply(mission)
-        return mission
-
-    def full_name(self) -> str:
-        return f"{self.site.name}{MAP_MISSION_DELIMITER}{self.name}"
-
-
-# CogsGuard vibes
-COGSGUARD_VIBES = [
-    Vibe("😐", "default"),
-    Vibe("❤️", "heart"),
-    Vibe("⚙️", "gear"),
-    Vibe("🌀", "scrambler"),
-    Vibe("🔗", "aligner"),
-    Vibe("⛏️", "miner"),
-    Vibe("🔭", "scout"),
+__all__ = [
+    "MAP_MISSION_DELIMITER",
+    "CoGameMission",
+    "CoGameMissionVariant",
+    "CoGameSite",
+    "CvCMission",
+    "NumCogsVariant",
 ]
 
 
-class Mission(MissionBase):
+class CvCMission(CoGameMission):
     """Mission configuration for CogsGuard game mode."""
 
-    # Agent configuration
-    cog: CogConfig = Field(default_factory=CogConfig)
-
-    wealth: int = Field(default=1)
+    max_steps: int = Field(default=10000)
 
-    # Collective initial resources
-    collective_initial_carbon: int = Field(default=10)
-    collective_initial_oxygen: int = Field(default=10)
-    collective_initial_germanium: int = Field(default=10)
-    collective_initial_silicon: int = Field(default=10)
-    collective_initial_heart: int = Field(default=5)
+    cog: CogConfig = Field(default_factory=lambda: CogConfig())
+    teams: dict[str, CogTeam] = Field(
+        default_factory=lambda: {
+            "cogs": CogTeam(name="cogs", num_agents=8, wealth=1),
+        }
+    )
 
-    # Clips Behavior - scramble cogs junctions to neutral
-    # Note: must start after initial_clips fires at timestep 10 (events fire alphabetically)
-    clips_scramble_start: int = Field(default=50)
-    clips_scramble_interval: int = Field(default=100)
-    clips_scramble_radius: int = Field(default=25)
+    clips: ClipsConfig = Field(default_factory=lambda: ClipsConfig())
 
-    # Clips Behavior - align neutral junctions to clips
-    clips_align_start: int = Field(default=100)
-    clips_align_interval: int = Field(default=100)
-    clips_align_radius: int = Field(default=25)
+    @property
+    def num_agents(self) -> int:
+        if self.num_cogs is not None:
+            return self.num_cogs
+        return sum(team.num_agents for team in self.teams.values())
 
-    # Station configs
-    wall: CvCWallConfig = Field(default_factory=CvCWallConfig)
+    def map_builder(self) -> AnyMapBuilderConfig:
+        """Return the map builder config. Override in subclasses for custom map generation."""
+        return self.site.map_builder
 
     def make_env(self) -> MettaGridConfig:
         """Create a MettaGridConfig from this mission.
@@ -192,129 +74,41 @@ class Mission(MissionBase):
         Returns:
             MettaGridConfig ready for environment creation
         """
-        map_builder = self.site.map_builder
-        num_cogs = self.num_cogs if self.num_cogs is not None else self.site.min_cogs
-
-        gear = GEAR
-        elements = ELEMENTS
-        resources_list = ["energy", "heart", "hp", "influence", *elements, *gear]
-        vibe_names = [vibe.name for vibe in COGSGUARD_VIBES]
-
-        extractor_objects = {
-            f"{resource}_extractor": SimpleExtractorConfig(resource=resource).station_cfg() for resource in elements
-        }
-        gear_objects = {f"{g}_station": GearStationConfig(gear_type=g).station_cfg() for g in gear}
-
-        # Create inventory observations for collective resources
-        collective_obs = [inv(f"collective.{resource}") for resource in elements]
 
         game = GameConfig(
-            map_builder=map_builder,
+            map_builder=self.map_builder(),
             max_steps=self.max_steps,
-            num_agents=num_cogs,
-            resource_names=resources_list,
-            vibe_names=vibe_names,
-            obs=ObsConfig(global_obs=GlobalObsConfig(obs=collective_obs, local_position=True)),
+            num_agents=self.num_agents,
+            resource_names=CvCConfig.RESOURCES,
+            vibe_names=CvCConfig.VIBE_NAMES,
+            obs=ObsConfig(
+                global_obs=GlobalObsConfig(
+                    obs=[inv(f"collective.{resource}") for resource in CvCConfig.ELEMENTS], local_position=True
+                )
+            ),
             actions=ActionsConfig(
-                move=MoveActionConfig(consumed_resources={"energy": self.cog.move_energy_cost}),
+                move=MoveActionConfig(consumed_resources=self.cog.action_cost),
                 noop=NoopActionConfig(),
-                change_vibe=ChangeVibeActionConfig(vibes=COGSGUARD_VIBES),
-            ),
-            agent=self.cog.agent_config(gear=gear, elements=elements).model_copy(
-                update={
-                    "rewards": {
-                        "aligned_junction_held": reward(
-                            game_stat("collective.aligned.junction.held"),
-                            weight=1.0 / self.max_steps,
-                            denoms=[numObjects("junction")],
-                        ),
-                    },
-                }
+                change_vibe=ChangeVibeActionConfig(vibes=CvCConfig.VIBES),
             ),
+            agent=self.cog.agent_config(team="cogs", max_steps=self.max_steps),
+            agents=[self.cog.agent_config(team="cogs", max_steps=self.max_steps) for _ in range(self.num_agents)],
             objects={
-                "wall": self.wall.station_cfg(),
-                "hub": HubConfig(map_name="hub", team="cogs").station_cfg(),
-                "junction": JunctionConfig(map_name="junction").station_cfg(),
-                "chest": CogsGuardChestConfig().station_cfg(),
-                **extractor_objects,
-                **gear_objects,
+                "wall": CvCWallConfig().station_cfg(),
+                "hub": CvCHubConfig().station_cfg(team="cogs"),
+                "junction": CvCJunctionConfig().station_cfg(),
+                "chest": CvCChestConfig().station_cfg(team="cogs"),
+                **{
+                    f"{resource}_extractor": CvCExtractorConfig(resource=resource).station_cfg()
+                    for resource in CvCConfig.ELEMENTS
+                },
+                **{f"{g}_station": CvCGearStationConfig(gear_type=g).station_cfg(team="cogs") for g in CvCConfig.GEAR},
             },
             collectives={
-                "cogs": CollectiveConfig(
-                    inventory=InventoryConfig(
-                        limits={
-                            "resources": ResourceLimitsConfig(min=10000, resources=elements),
-                            "hearts": ResourceLimitsConfig(min=65535, resources=["heart"]),
-                        },
-                        initial={
-                            "carbon": self.collective_initial_carbon * self.wealth,
-                            "oxygen": self.collective_initial_oxygen * self.wealth,
-                            "germanium": self.collective_initial_germanium * self.wealth,
-                            "silicon": self.collective_initial_silicon * self.wealth,
-                            "heart": self.collective_initial_heart * self.wealth,
-                        },
-                    ),
-                ),
-                "clips": CollectiveConfig(),
-            },
-            events={
-                "initial_clips": EventConfig(
-                    name="initial_clips",
-                    target_tag=typeTag("junction"),
-                    timesteps=[10],
-                    mutations=[alignTo("clips")],
-                    max_targets=1,
-                ),
-                "cogs_to_neutral": EventConfig(
-                    name="cogs_to_neutral",
-                    target_tag=typeTag("junction"),
-                    timesteps=periodic(
-                        start=self.clips_scramble_start,
-                        period=self.clips_scramble_interval,
-                        end_period=self.clips_scramble_interval // 5,
-                        end=self.max_steps,
-                    ),
-                    filters=[
-                        isNear(typeTag("junction"), [isAlignedTo("clips")], radius=self.clips_scramble_radius),
-                        isAlignedTo("cogs"),
-                    ],
-                    mutations=[alignTo(None)],
-                    max_targets=1,
-                ),
-                "neutral_to_clips": EventConfig(
-                    name="neutral_to_clips",
-                    target_tag=typeTag("junction"),
-                    timesteps=periodic(
-                        start=self.clips_align_start,
-                        period=self.clips_align_interval,
-                        end_period=self.clips_align_interval // 5,
-                        end=self.max_steps,
-                    ),
-                    filters=[
-                        isNear(typeTag("junction"), [isAlignedTo("clips")], radius=self.clips_align_radius),
-                        isAlignedTo(None),
-                    ],
-                    mutations=[alignTo("clips")],
-                    max_targets=1,
-                    fallback="cogs_to_neutral",
-                ),
-                # If the Clips can't find any junctions near them, align a random junction
-                "presence_check": EventConfig(
-                    name="presence_check",
-                    target_tag=typeTag("junction"),
-                    timesteps=periodic(
-                        start=self.clips_scramble_start,
-                        period=self.clips_scramble_interval,
-                        end=self.max_steps,
-                    ),
-                    filters=[
-                        isNear(typeTag("junction"), [isAlignedTo("clips")], radius=self.clips_scramble_radius),
-                    ],
-                    mutations=[],
-                    max_targets=1,
-                    fallback="initial_clips",
-                ),
+                **{team.name: team.collective_config() for team in self.teams.values()},
+                "clips": self.clips.collective_config(),
             },
+            events=self.clips.events(max_steps=self.max_steps),
         )
 
         env = MettaGridConfig(game=game)
@@ -328,9 +122,3 @@ class Mission(MissionBase):
             env.label += f".{variant.name}"
 
         return env
-
-
-# Backwards compatibility alias
-CogsGuardMission = Mission
-
-AnyMission = Mission