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

cogames/cogs_vs_clips/cogsguard_reward_variants.py

@@ -0,0 +1,153 @@
+"""Reward preset wiring for the CogsGuard (Cogs vs Clips) mission.
+
+The mission has a single "true" objective signal, plus optional shaping variants.
+Reward variants are stackable; each one adds additional shaping signals on top of the
+mission's default objective rewards.
+"""
+
+from __future__ import annotations
+
+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"]
+
+AVAILABLE_REWARD_VARIANTS: tuple[CogsGuardRewardVariant, ...] = (
+    "objective",
+    "no_objective",
+    "milestones",
+    "credit",
+    "penalize_vibe_change",
+)
+
+_OBJECTIVE_STAT_KEY = "aligned_junction_held"
+
+
+def _apply_milestones(rewards: dict[str, AgentReward], *, max_junctions: int = 100) -> None:
+    """Add milestone shaping rewards onto an existing baseline.
+
+    Args:
+        rewards: Rewards dict to modify in-place.
+        max_junctions: Maximum expected number of junctions for capping rewards.
+            Defaults to 100 as a reasonable upper bound for most maps.
+    """
+    w_junction_aligned = 1.0
+    w_scramble_act = 0.5
+    w_align_act = 1.0
+
+    # Max caps based on expected junction counts
+    max_junction_aligned = w_junction_aligned * max_junctions
+    max_scramble = w_scramble_act * max_junctions
+    max_align = w_align_act * max_junctions
+
+    rewards["aligned_junctions"] = reward(
+        stat("collective.junction"),
+        weight=w_junction_aligned,
+        max=max_junction_aligned,
+    )
+
+    rewards["junction_scrambled_by_agent"] = reward(
+        stat("junction.scrambled_by_agent"),
+        weight=w_scramble_act,
+        max=max_scramble,
+    )
+    rewards["junction_aligned_by_agent"] = reward(
+        stat("junction.aligned_by_agent"),
+        weight=w_align_act,
+        max=max_align,
+    )
+
+
+def _apply_penalize_vibe_change(rewards: dict[str, AgentReward]) -> None:
+    """Add penalty for vibe changes to discourage spamming."""
+    w_vibe_change = -0.01
+    rewards["vibe_change_penalty"] = reward(stat("action.change_vibe.success"), weight=w_vibe_change)
+
+
+def _apply_credit(rewards: dict[str, AgentReward]) -> None:
+    """Add dense precursor shaping rewards onto an existing baseline."""
+    w_heart = 0.05
+    cap_heart = 0.5
+    w_align_gear = 0.2
+    cap_align_gear = 0.4
+    w_scramble_gear = 0.2
+    cap_scramble_gear = 0.4
+    w_element_gain = 0.001
+    cap_element_gain = 0.1
+
+    # Stats rewards for gains as a single map
+    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),
+        "scrambler_gained": reward(stat("scrambler.gained"), 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),
+        "silicon_gained": reward(stat("silicon.gained"), weight=w_element_gain, max=cap_element_gain),
+    }
+    rewards.update(gain_rewards)
+
+    # Collective deposit rewards
+    w_deposit = 0.002
+    cap_deposit = 0.2
+    deposit_rewards: dict[str, AgentReward] = {
+        f"collective_{element}_deposited": reward(
+            stat(f"collective.{element}.deposited"), weight=w_deposit, max=cap_deposit
+        )
+        for element in ["carbon", "oxygen", "germanium", "silicon"]
+    }
+    rewards.update(deposit_rewards)
+
+
+def apply_reward_variants(env: MettaGridConfig, *, variants: str | Sequence[str] | None = None) -> None:
+    """Apply CogsGuard reward variants to `env`.
+
+    Variants are stackable:
+    - `objective`: no-op marker; keeps the mission's default objective reward wiring.
+    - `no_objective`: disables the objective stat reward (`junction.held`).
+    - `milestones`: adds shaped rewards for aligning/scrambling junctions and holding more junctions.
+    - `credit`: adds additional dense shaping for precursor behaviors (resources/gear/deposits).
+    - `penalize_vibe_change`: adds a penalty for vibe changes to discourage spamming.
+    """
+    if not variants:
+        return
+
+    variant_names = [variants] if isinstance(variants, str) else list(variants)
+
+    reward_variants: list[CogsGuardRewardVariant] = []
+    for variant_name in variant_names:
+        if variant_name not in AVAILABLE_REWARD_VARIANTS:
+            available = ", ".join(AVAILABLE_REWARD_VARIANTS)
+            raise ValueError(f"Unknown Cogsguard reward variant '{variant_name}'. Available: {available}")
+        variant = cast(CogsGuardRewardVariant, variant_name)
+        if variant in reward_variants:
+            continue
+        reward_variants.append(variant)
+
+    enabled = set(reward_variants)
+    if enabled <= {"objective"}:
+        return
+
+    # Start from the mission's existing objective baseline to preserve its scaling.
+    rewards = dict(env.game.agent.rewards)
+
+    if "no_objective" in enabled:
+        rewards.pop(_OBJECTIVE_STAT_KEY, None)
+    if "milestones" in enabled:
+        _apply_milestones(rewards)
+    if "credit" in enabled:
+        _apply_credit(rewards)
+    if "penalize_vibe_change" in enabled:
+        _apply_penalize_vibe_change(rewards)
+
+    env.game.agent.rewards = rewards
+
+    # Deterministic label suffix order (exclude "objective").
+    for variant in AVAILABLE_REWARD_VARIANTS:
+        if variant == "objective":
+            continue
+        if variant in enabled:
+            env.label += f".{variant}"

cogames/cogs_vs_clips/cogsguard_tutorial.py

@@ -0,0 +1,56 @@
+"""CogsGuard tutorial mission configuration."""
+
+from cogames.cogs_vs_clips.mission import Mission, Site
+from cogames.cogs_vs_clips.procedural import MachinaArena
+from mettagrid.mapgen.mapgen import MapGen
+from mettagrid.mapgen.scenes.base_hub import BaseHubConfig
+
+
+def make_cogsguard_tutorial_site() -> Site:
+    """Create a smaller, simpler CogsGuard arena for the tutorial."""
+    hub_config = BaseHubConfig(
+        corner_bundle="extractors",
+        cross_bundle="none",
+        cross_distance=5,
+        hub_width=15,
+        hub_height=15,
+        outer_clearance=2,
+        stations=[
+            "aligner_station",
+            "scrambler_station",
+            "miner_station",
+            "scout_station",
+            "chest",
+        ],
+    )
+    map_builder = MapGen.Config(
+        width=35,
+        height=35,
+        instance=MachinaArena.Config(
+            spawn_count=1,
+            building_coverage=0.05,
+            hub=hub_config,
+        ),
+    )
+    return Site(
+        name="cogsguard_tutorial",
+        description="CogsGuard tutorial arena - small map for learning",
+        map_builder=map_builder,
+        min_cogs=1,
+        max_cogs=1,
+    )
+
+
+CogsGuardTutorialMission = Mission(
+    name="tutorial",
+    description="Learn the basics of CogsGuard: Roles, Resources, and Territory Control.",
+    site=make_cogsguard_tutorial_site(),
+    num_cogs=1,
+    max_steps=2000,
+    # Generous initial resources for learning
+    collective_initial_carbon=50,
+    collective_initial_oxygen=50,
+    collective_initial_germanium=50,
+    collective_initial_silicon=50,
+    collective_initial_heart=10,
+)

cogames/cogs_vs_clips/evals/README.md

@@ -19,9 +19,8 @@ Diagnostic missions test specific skills in isolation with controlled, repeatabl
 #### Navigation & Delivery
 
 - `diagnostic_chest_navigation1/2/3` - Navigate to chest and deposit hearts (varying difficulty)
-- `diagnostic_chest_near` - Chest nearby, test deposit mechanics
-- `diagnostic_chest_search` - Find chest through exploration
-- `diagnostic_chest_deposit_near/search` - Combined navigation and deposit tests
+- `diagnostic_chest_deposit_near` - Chest nearby, test deposit mechanics
+- `diagnostic_chest_deposit_search` - Find chest through exploration
 
 #### Resource Extraction
 
@@ -32,18 +31,13 @@ Diagnostic missions test specific skills in isolation with controlled, repeatabl
 
 #### Assembly
 
-- `diagnostic_assembler_near` - Assemble hearts at nearby assembler
-- `diagnostic_assembler_search` - Find assembler and craft hearts
+- `diagnostic_assemble_seeded_near` - Agents pre-seeded with resources, assemble hearts at nearby hub
+- `diagnostic_assemble_seeded_search` - Agents pre-seeded, find hub and craft hearts
 
 #### Energy Management
 
 - `diagnostic_charge_up` - Test charging mechanics and energy management
 
-#### Unclipping
-
-- `diagnostic_unclip_craft` - Craft unclip items to restore clipped extractors
-- `diagnostic_unclip_preseed` - Unclip with pre-seeded inventory
-
 #### Complex Scenarios
 
 - `diagnostic_radial` - Radial resource layout with chorus assembly
@@ -68,7 +62,7 @@ uv run cogames play --mission evals.diagnostic_extract_missing_oxygen --cogs 2
 uv run cogames play --mission evals.diagnostic_radial_hard --cogs 1
 
 # With policy
-uv run cogames play --mission evals.diagnostic_unclip_craft -p scripted_baseline --cogs 1
+uv run cogames play --mission evals.diagnostic_radial -p baseline --cogs 1
 ```
 
 ---
@@ -100,7 +94,7 @@ uv run cogames play --mission hello_world.oxygen_bottleneck --cogs 2
 
 #### energy_starved
 
-**Challenge:** Low energy regen and weak chargers require careful energy management.
+**Challenge:** Low energy regen and weak junctions require careful energy management.
 
 **Variants Applied:**
 
@@ -192,7 +186,7 @@ uv run cogames play --mission hello_world.oxygen_bottleneck --cogs 2 --variant c
 uv run cogames play --mission hello_world.energy_starved --cogs 2 --variant compass --variant small_50
 
 # With policy
-uv run cogames play --mission hello_world.single_use_swarm --cogs 4 -p scripted_baseline
+uv run cogames play --mission hello_world.single_use_swarm --cogs 4 -p baseline
 ```
 
 ---
@@ -206,16 +200,16 @@ For systematic evaluation across multiple missions and configurations:
 ```bash
 # Evaluate on integrated eval suite
 uv run python packages/cogames/scripts/run_evaluation.py \
-  --agent cogames.policy.nim_agents.agents.ThinkyAgentsMultiPolicy \
+  --policy thinky \
   --mission-set integrated_evals \
   --cogs 4 \
   --repeats 2
 
 # Evaluate specific agent
 uv run python packages/cogames/scripts/run_evaluation.py \
-  --agent simple \
+  --policy baseline \
   --steps 1000 \
-  --output eval_simple.json
+  --output eval_baseline.json
 ```
 
 ### Using in Curriculum Training

cogames/cogs_vs_clips/evals/cogsguard_evals.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from cogames.cogs_vs_clips.mission import CogsGuardMission, Site
+from cogames.cogs_vs_clips.mission_utils import get_map
+
+MAPS_DIR = Path(__file__).resolve().parent.parent.parent / "maps"
+
+COGSGUARD_EVALS_BASE = Site(
+    name="cogsguard_evals",
+    description="CogsGuard evaluation arenas.",
+    map_builder=get_map("evals/eval_balanced_spread.map"),
+    min_cogs=1,
+    max_cogs=20,
+)
+
+
+def _count_spawn_pads(map_path: Path) -> int:
+    text = map_path.read_text()
+    if "map_data:" not in text:
+        raise ValueError(f"Missing map_data block in {map_path}")
+    map_section = text.split("map_data:", 1)[1].split("char_to_map_name:", 1)[0]
+    count = map_section.count("@")
+    if count <= 0:
+        raise ValueError(f"No spawn pads found in {map_path}")
+    return count
+
+
+def _make_eval_site(map_name: str, num_cogs: int) -> Site:
+    site = COGSGUARD_EVALS_BASE.model_copy(
+        update={
+            "map_builder": get_map(map_name),
+            "min_cogs": num_cogs,
+            "max_cogs": num_cogs,
+        }
+    )
+    return site
+
+
+def _description_from_stem(stem: str) -> str:
+    display = stem
+    if display.startswith("eval_"):
+        display = display[len("eval_") :]
+    display = display.replace("_", " ")
+    return f"CogsGuard eval: {display}."
+
+
+COGSGUARD_EVAL_MAPS: list[str] = [
+    "evals/eval_balanced_spread.map",
+    "evals/eval_clip_oxygen.map",
+    "evals/eval_collect_resources.map",
+    "evals/eval_collect_resources_medium.map",
+    "evals/eval_collect_resources_hard.map",
+    "evals/eval_divide_and_conquer.map",
+    "evals/eval_energy_starved.map",
+    "evals/eval_multi_coordinated_collect_hard.map",
+    "evals/eval_oxygen_bottleneck.map",
+    "evals/eval_single_use_world.map",
+    "evals/extractor_hub_30x30.map",
+    "evals/extractor_hub_50x50.map",
+    "evals/extractor_hub_70x70.map",
+    "evals/extractor_hub_80x80.map",
+    "evals/extractor_hub_100x100.map",
+]
+
+COGSGUARD_EVAL_COGS = {map_name: _count_spawn_pads(MAPS_DIR / map_name) for map_name in COGSGUARD_EVAL_MAPS}
+
+COGSGUARD_EVAL_MISSIONS: list[CogsGuardMission] = []
+for map_name in COGSGUARD_EVAL_MAPS:
+    stem = Path(map_name).stem
+    num_cogs = COGSGUARD_EVAL_COGS[map_name]
+    site = _make_eval_site(map_name, num_cogs)
+    COGSGUARD_EVAL_MISSIONS.append(
+        CogsGuardMission(
+            name=stem,
+            description=_description_from_stem(stem),
+            site=site,
+            num_cogs=num_cogs,
+        )
+    )