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

cogames/cogs_vs_clips/mission.py

@@ -1,188 +1,73 @@
 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.cogs_vs_clips.weather import WeatherConfig
+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 CollectiveConfig, 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)
+    total_junctions: int = Field(default=118, description="Total junctions on the map (for curriculum scaling)")
 
-    # 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())
+    weather: WeatherConfig = Field(default_factory=lambda: WeatherConfig())
 
-    # 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 +77,57 @@ 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]
-
+        team_objs = list(self.teams.values())
         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),
             ),
+            agents=[
+                self.cog.agent_config(team=t.name, max_steps=self.max_steps)
+                for t in team_objs
+                for _ in range(t.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(),
+                "junction": CvCJunctionConfig().station_cfg(),
+                **{
+                    f"{resource}_extractor": CvCExtractorConfig(resource=resource).station_cfg()
+                    for resource in CvCConfig.ELEMENTS
+                },
+                **{
+                    f"{t.short_name}:hub": CvCHubConfig().station_cfg(team=t.short_name, collective=t.name)
+                    for t in team_objs
+                },
+                **{
+                    f"{t.short_name}:chest": CvCChestConfig().station_cfg(team=t.short_name, collective=t.name)
+                    for t in team_objs
+                },
+                **{
+                    f"{t.short_name}:{g}": CvCGearStationConfig(gear_type=g).station_cfg(
+                        team=t.short_name, collective=t.name
+                    )
+                    for t in team_objs
+                    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",
-                ),
+                **{t.name: t.collective_config() for t in team_objs},
+                **self.clips.collectives(),
+                "neutral": CollectiveConfig(name="neutral"),
             },
+            events=self._merge_events(),
         )
 
         env = MettaGridConfig(game=game)
@@ -329,8 +142,11 @@ class Mission(MissionBase):
 
         return env
 
-
-# Backwards compatibility alias
-CogsGuardMission = Mission
-
-AnyMission = Mission
+    def _merge_events(self) -> dict:
+        """Merge clips and weather events, raising on key conflicts."""
+        clips_events = self.clips.events(max_steps=self.max_steps)
+        weather_events = self.weather.events(max_steps=self.max_steps)
+        overlap = set(clips_events) & set(weather_events)
+        if overlap:
+            raise ValueError(f"Overlapping event keys between clips and weather: {overlap}")
+        return {**clips_events, **weather_events}

cogames/cogs_vs_clips/missions.py

@@ -1,6 +1,6 @@
 from functools import lru_cache
 
-from cogames.cogs_vs_clips.mission import AnyMission, Mission
+from cogames.cogs_vs_clips.mission import CvCMission
 from cogames.cogs_vs_clips.sites import (
     COGSGUARD_ARENA,
     COGSGUARD_MACHINA_1,
@@ -11,9 +11,9 @@ from mettagrid.config.mettagrid_config import MettaGridConfig
 # CogsGuard Missions
 
 
-def make_cogsguard_mission(num_agents: int = 10, max_steps: int = 10000) -> Mission:
+def make_cogsguard_mission(num_agents: int = 10, max_steps: int = 10000) -> CvCMission:
     """Create a CogsGuard mission with configurable parameters (Machina1 layout)."""
-    return Mission(
+    return CvCMission(
         name="basic",
         description="Basic CogsGuard mission (Machina1 layout)",
         site=make_cogsguard_machina1_site(num_agents),
@@ -22,7 +22,7 @@ def make_cogsguard_mission(num_agents: int = 10, max_steps: int = 10000) -> Miss
     )
 
 
-CogsGuardMachina1Mission = Mission(
+CogsGuardMachina1Mission = CvCMission(
     name="basic",
     description="CogsGuard Machina1 - compete to control junctions with gear abilities.",
     site=COGSGUARD_MACHINA_1,
@@ -30,7 +30,7 @@ CogsGuardMachina1Mission = Mission(
     max_steps=10000,
 )
 
-CogsGuardBasicMission = Mission(
+CogsGuardBasicMission = CvCMission(
     name="basic",
     description="CogsGuard Arena - compact training map with gear abilities.",
     site=COGSGUARD_ARENA,
@@ -39,48 +39,38 @@ CogsGuardBasicMission = Mission(
 )
 
 
-_CORE_MISSIONS: list[AnyMission] = [
+_CORE_MISSIONS: list[CvCMission] = [
     CogsGuardMachina1Mission,
     CogsGuardBasicMission,
 ]
 
 
-def get_core_missions() -> list[AnyMission]:
+def get_core_missions() -> list[CvCMission]:
     return list(_CORE_MISSIONS)
 
 
-def get_legacy_missions() -> list[Mission]:
-    """Get legacy (pre-CogsGuard) missions for backward compatibility.
-
-    Legacy missions have been removed. Returns empty list.
-    """
-    return []
-
-
-def _build_eval_missions() -> list[AnyMission]:
-    from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS
-    from cogames.cogs_vs_clips.evals.integrated_evals import EVAL_MISSIONS as INTEGRATED_EVAL_MISSIONS
+def _build_eval_missions() -> list[CvCMission]:
+    from cogames.cogs_vs_clips.evals.integrated_evals import EVAL_MISSIONS as INTEGRATED_EVAL_MISSIONS  # noqa: PLC0415
 
     return [
-        *COGSGUARD_EVAL_MISSIONS,
         *INTEGRATED_EVAL_MISSIONS,
     ]
 
 
 @lru_cache(maxsize=1)
-def get_missions() -> list[AnyMission]:
+def get_missions() -> list[CvCMission]:
     return [*_CORE_MISSIONS, *_build_eval_missions()]
 
 
-def __getattr__(name: str) -> list[AnyMission]:
+def make_game(num_cogs: int = 2, map_name: str = "training_facility_open_1.map") -> MettaGridConfig:
+    """Create a default CogsGuard game configuration."""
+    mission = make_cogsguard_mission(num_agents=num_cogs)
+    return mission.make_env()
+
+
+def __getattr__(name: str) -> list[CvCMission]:
     if name == "MISSIONS":
         missions = get_missions()
         globals()["MISSIONS"] = missions
         return missions
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-
-
-def make_game(num_cogs: int = 2, map_name: str = "training_facility_open_1.map") -> MettaGridConfig:
-    """Create a default CogsGuard game configuration."""
-    mission = make_cogsguard_mission(num_agents=num_cogs)
-    return mission.make_env()

cogames/cogs_vs_clips/{cogsguard_reward_variants.py → reward_variants.py}

@@ -7,13 +7,20 @@ mission's default objective rewards.
 
 from __future__ import annotations
 
+import json
 from typing import Literal, Sequence, cast
 
 from mettagrid.config.game_value import stat
 from mettagrid.config.mettagrid_config import MettaGridConfig
 from mettagrid.config.reward_config import AgentReward, reward
 
-CogsGuardRewardVariant = Literal["credit", "milestones", "no_objective", "penalize_vibe_change", "objective"]
+CogsGuardRewardVariant = Literal[
+    "credit",
+    "milestones",
+    "no_objective",
+    "penalize_vibe_change",
+    "objective",
+]
 
 AVAILABLE_REWARD_VARIANTS: tuple[CogsGuardRewardVariant, ...] = (
     "objective",
@@ -82,7 +89,9 @@ def _apply_credit(rewards: dict[str, AgentReward]) -> None:
     gain_rewards: dict[str, AgentReward] = {
         "heart_gained": reward(stat("heart.gained"), weight=w_heart, max=cap_heart),
         "aligner_gained": reward(stat("aligner.gained"), weight=w_align_gear, max=cap_align_gear),
+        "aligner_lost": reward(stat("aligner.lost"), weight=-w_align_gear, max=-cap_align_gear),
         "scrambler_gained": reward(stat("scrambler.gained"), weight=w_scramble_gear, max=cap_scramble_gear),
+        "scrambler_lost": reward(stat("scrambler.lost"), weight=-w_scramble_gear, max=-cap_scramble_gear),
         "carbon_gained": reward(stat("carbon.gained"), weight=w_element_gain, max=cap_element_gain),
         "oxygen_gained": reward(stat("oxygen.gained"), weight=w_element_gain, max=cap_element_gain),
         "germanium_gained": reward(stat("germanium.gained"), weight=w_element_gain, max=cap_element_gain),
@@ -115,7 +124,18 @@ def apply_reward_variants(env: MettaGridConfig, *, variants: str | Sequence[str]
     if not variants:
         return
 
-    variant_names = [variants] if isinstance(variants, str) else list(variants)
+    # Parse JSON-encoded list strings (e.g., '["milestones"]' from sweeps)
+    if isinstance(variants, str):
+        if variants.startswith("["):
+            try:
+                parsed = json.loads(variants)
+                variant_names = list(parsed) if isinstance(parsed, list) else [variants]
+            except json.JSONDecodeError:
+                variant_names = [variants]
+        else:
+            variant_names = [variants]
+    else:
+        variant_names = list(variants)
 
     reward_variants: list[CogsGuardRewardVariant] = []
     for variant_name in variant_names: