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

cogames/play.py

@@ -1,21 +1,160 @@
 """Game playing functionality for CoGames."""
 
 import logging
+import uuid
 from pathlib import Path
 from typing import Optional
 
+from rich import box
 from rich.console import Console
+from rich.table import Table
 
+from metta_alo.pure_single_episode_runner import PureSingleEpisodeResult
+from metta_alo.rollout import run_single_episode
 from mettagrid import MettaGridConfig
-from mettagrid.policy.loader import initialize_or_load_policy
 from mettagrid.policy.policy import PolicySpec
-from mettagrid.policy.policy_env_interface import PolicyEnvInterface
 from mettagrid.renderer.renderer import RenderMode
-from mettagrid.simulator.replay_log_writer import ReplayLogWriter
-from mettagrid.simulator.rollout import Rollout
 
 logger = logging.getLogger("cogames.play")
 
+# Resources and gear types for CogsGuard
+ELEMENTS = ["carbon", "oxygen", "germanium", "silicon"]
+GEAR = ["miner", "aligner", "scrambler", "scout"]
+
+
+def _print_episode_stats(console: Console, results: PureSingleEpisodeResult) -> None:
+    """Print episode statistics in a formatted table."""
+    stats = results.stats
+    total_reward = sum(results.rewards)
+
+    # Aggregate agent stats
+    agent_stats = stats.get("agent", [])
+    totals: dict[str, float] = {}
+    for agent in agent_stats:
+        for key, value in agent.items():
+            totals[key] = totals.get(key, 0) + value
+
+    # Check if this is a CogsGuard mission (has collective stats)
+    collective_stats = stats.get("collective", {})
+    cogs_stats = collective_stats.get("cogs", {})
+    clips_stats = collective_stats.get("clips", {})
+
+    if cogs_stats or clips_stats:
+        # CogsGuard mission - show relevant stats
+        _print_cogsguard_stats(console, totals, cogs_stats, clips_stats, total_reward)
+    else:
+        # Standard mission - show basic stats
+        _print_standard_stats(console, totals, total_reward)
+
+
+def _print_cogsguard_stats(
+    console: Console,
+    agent_totals: dict[str, float],
+    cogs_stats: dict[str, float],
+    clips_stats: dict[str, float],
+    total_reward: float,
+) -> None:
+    """Print CogsGuard-specific statistics."""
+    table = Table(title="Episode Stats", box=box.ROUNDED, show_header=True, header_style="bold cyan")
+    table.add_column("Category", style="yellow")
+    table.add_column("Stat", style="white")
+    table.add_column("Gained", style="green", justify="right")
+    table.add_column("Lost", style="red", justify="right")
+    table.add_column("Final", style="cyan", justify="right")
+
+    sections_added = 0
+
+    # Junctions: gained (aligned) | lost | final (current count)
+    junctions_added = False
+    for team, team_stats, color in [("cogs", cogs_stats, "green"), ("clips", clips_stats, "red")]:
+        gained = int(team_stats.get("junction.gained", 0))
+        lost = int(team_stats.get("junction.lost", 0))
+        final = int(team_stats.get("junction", 0))
+        if gained > 0 or lost > 0 or final > 0:
+            if not junctions_added:
+                table.add_row("[bold]Junctions[/bold]", "", "", "", "")
+                junctions_added = True
+                sections_added += 1
+            table.add_row(
+                "",
+                f"[{color}]{team}[/{color}]",
+                f"[{color}]{gained}[/{color}]",
+                f"[{color}]{lost}[/{color}]",
+                f"[{color}]{final}[/{color}]",
+            )
+
+    # Gear: gained | lost | final (net = gained - lost)
+    gear_added = False
+    for gear in GEAR:
+        gained = int(agent_totals.get(f"{gear}.gained", 0))
+        lost = int(agent_totals.get(f"{gear}.lost", 0))
+        final = gained - lost
+        if gained > 0 or lost > 0:
+            if not gear_added:
+                if sections_added > 0:
+                    table.add_section()
+                table.add_row("[bold]Gear[/bold]", "", "", "", "")
+                gear_added = True
+                sections_added += 1
+            table.add_row("", gear, str(gained), str(lost), str(final))
+
+    # Hearts (in Gear section)
+    hearts_gained = int(agent_totals.get("heart.gained", 0))
+    hearts_lost = int(agent_totals.get("heart.lost", 0))
+    if hearts_gained > 0 or hearts_lost > 0:
+        if not gear_added:
+            if sections_added > 0:
+                table.add_section()
+            table.add_row("[bold]Gear[/bold]", "", "", "", "")
+            gear_added = True
+            sections_added += 1
+        table.add_row("", "hearts", str(hearts_gained), str(hearts_lost), "")
+
+    # Resources: gained (deposited) | lost (withdrawn) | final (current amount)
+    resources_added = False
+    for resource in ELEMENTS:
+        gained = int(cogs_stats.get(f"collective.{resource}.deposited", 0))
+        lost = int(cogs_stats.get(f"collective.{resource}.withdrawn", 0))
+        final = int(cogs_stats.get(f"collective.{resource}.amount", 0))
+        if gained > 0 or lost > 0 or final > 0:
+            if not resources_added:
+                if sections_added > 0:
+                    table.add_section()
+                table.add_row("[bold]Resources[/bold]", "", "", "", "")
+                resources_added = True
+                sections_added += 1
+            table.add_row("", resource, str(gained), str(lost), str(final))
+
+    # Total reward at bottom
+    if sections_added > 0:
+        table.add_section()
+    table.add_row("[bold]Reward[/bold]", "total", f"{total_reward:.2f}", "", "")
+
+    console.print(table)
+
+
+def _print_standard_stats(console: Console, agent_totals: dict[str, float], total_reward: float) -> None:
+    """Print standard statistics for non-CogsGuard missions."""
+    # Filter for interesting stats
+    interesting = {}
+    for key, value in agent_totals.items():
+        if value != 0 and any(pattern in key for pattern in [".gained", ".lost", ".deposited", ".withdrawn", "heart"]):
+            interesting[key] = value
+
+    table = Table(title="Episode Stats", box=box.ROUNDED, show_header=True, header_style="bold cyan")
+    table.add_column("Stat", style="white")
+    table.add_column("Value", style="green", justify="right")
+
+    for key in sorted(interesting.keys()):
+        table.add_row(key, f"{int(interesting[key])}")
+
+    # Total reward at bottom
+    if interesting:
+        table.add_section()
+    table.add_row("[bold]Reward (total)[/bold]", f"{total_reward:.2f}")
+
+    console.print(table)
+
 
 def play(
     console: Console,
@@ -34,48 +173,42 @@ def play(
         policy_spec: Policy specification (class path and optional data path)
         game_name: Human-readable name of the game (used for logging/metadata)
         seed: Random seed
-        render_mode: Render mode - "gui", "unicode", or "none"
+        render_mode: Render mode - "gui", "vibescope", "unicode", or "none"
         save_replay: Optional directory path to save replay. Directory will be created if it doesn't exist.
             Replay will be saved with a unique UUID-based filename.
     """
 
     logger.debug("Starting play session", extra={"game_name": game_name})
 
-    policy_env_info = PolicyEnvInterface.from_mg_cfg(env_cfg)
-    policy = initialize_or_load_policy(policy_env_info, policy_spec)
-    agent_policies = [policy.agent_policy(agent_id) for agent_id in range(env_cfg.game.num_agents)]
-
-    # Set up replay writer if requested
-    event_handlers = []
-    replay_writer = None
+    replay_path = None
     if save_replay:
-        replay_writer = ReplayLogWriter(str(save_replay))
-        event_handlers.append(replay_writer)
-
-    # Create simulator and renderer
-    rollout = Rollout(
-        env_cfg,
-        agent_policies,
-        render_mode=render_mode,
-        seed=seed,
-        event_handlers=event_handlers,
-    )
+        save_replay.mkdir(parents=True, exist_ok=True)
+        replay_path = save_replay / f"{uuid.uuid4()}.json.z"
+
     try:
-        rollout.run_until_done()
+        results, _replay = run_single_episode(
+            policy_specs=[policy_spec],
+            assignments=[0] * env_cfg.game.num_agents,
+            env=env_cfg,
+            results_uri=None,
+            replay_uri=str(replay_path) if replay_path else None,
+            seed=seed,
+            device="cpu",
+            render_mode=render_mode,
+        )
     except KeyboardInterrupt:
         logger.info("Interrupted; ending episode early.")
-        rollout._sim.end_episode()
         return
 
     # Print summary
     console.print("\n[bold green]Episode Complete![/bold green]")
-    console.print(f"Steps: {rollout._sim.current_step}")
-    console.print(f"Total Rewards: {rollout._sim.episode_rewards}")
-    console.print(f"Final Reward Sum: {float(sum(rollout._sim.episode_rewards)):.2f}")
+    console.print(f"Steps: {results.steps}")
+
+    # Print episode stats
+    _print_episode_stats(console, results)
 
     # Print replay command if replay was saved
-    if replay_writer:
-        for replay_path in replay_writer.get_written_replay_paths():
-            console.print("\n[bold cyan]Replay saved![/bold cyan]")
-            console.print("To watch the replay, run:")
-            console.print(f"[bold green]cogames replay {replay_path}[/bold green]")
+    if replay_path:
+        console.print("\n[bold cyan]Replay saved![/bold cyan]")
+        console.print("To watch the replay, run:")
+        console.print(f"[bold green]cogames replay {replay_path}[/bold green]")

cogames/policy/chaos_monkey.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+import random
+from typing import Optional
+
+from mettagrid.policy.policy import AgentPolicy, MultiAgentPolicy
+from mettagrid.policy.policy_env_interface import PolicyEnvInterface
+from mettagrid.simulator import Action, AgentObservation, Simulation
+
+
+class _ChaosMonkeyAgent(AgentPolicy):
+    def __init__(self, policy_env_info: PolicyEnvInterface, fail_step: int, fail_probability: float):
+        super().__init__(policy_env_info)
+        self._fail_step = fail_step
+        self._fail_probability = fail_probability
+        self._step = 0
+        self._failed = False
+        self._noop = Action(name="noop")
+
+    def reset(self, simulation: Optional[Simulation] = None) -> None:
+        self._step = 0
+        self._failed = False
+
+    def step(self, obs: AgentObservation) -> Action:
+        if self._failed:
+            return self._noop
+
+        if self._step >= self._fail_step and random.random() < self._fail_probability:
+            self._failed = True
+            raise RuntimeError(f"Chaos monkey triggered at step {self._step}")
+
+        self._step += 1
+        return self._noop
+
+
+class ChaosMonkeyPolicy(MultiAgentPolicy):
+    """A scripted policy that intentionally fails mid-episode to test robustness."""
+
+    short_names = ["chaos-monkey"]
+
+    def __init__(
+        self,
+        policy_env_info: PolicyEnvInterface,
+        fail_step: int | str = 10,
+        fail_probability: float | int | str = 1.0,
+        device: str = "cpu",
+        **_: object,
+    ):
+        super().__init__(policy_env_info, device=device)
+        self._fail_step = int(fail_step)
+        self._fail_probability = float(fail_probability)
+
+    def agent_policy(self, agent_id: int) -> AgentPolicy:
+        return _ChaosMonkeyAgent(self.policy_env_info, self._fail_step, self._fail_probability)

cogames/policy/nim_agents/__init__.py

@@ -1,18 +1,35 @@
-"""Nim-based agent policies for CoGames."""
+"""Legacy import shim for Nim-based agents.
 
-from cogames.policy.nim_agents import agents  # noqa: F401
+Historically, Nim agents lived at `cogames.policy.nim_agents`. They were moved
+to the optional `cogames-agents` package under `cogames_agents.policy.nim_agents`.
+
+Old policy bundles may still import the legacy path; this module preserves
+backwards compatibility by re-exporting the new implementation when available.
+"""
+
+from __future__ import annotations
+
+try:
+    import cogames.policy.nim_agents.agents as agents  # noqa: F401
+    from cogames.policy.nim_agents.agents import (  # noqa: F401
+        CogsguardAlignAllAgentsMultiPolicy,
+        LadyBugAgentsMultiPolicy,
+        RaceCarAgentsMultiPolicy,
+        RandomAgentsMultiPolicy,
+        ThinkyAgentsMultiPolicy,
+    )
+except ModuleNotFoundError as exc:
+    if exc.name and (exc.name == "cogames_agents" or exc.name.startswith("cogames_agents.")):
+        raise ModuleNotFoundError(
+            "Legacy import `cogames.policy.nim_agents` requires optional dependency "
+            "`cogames-agents` (install `cogames[agents]`)."
+        ) from exc
+    raise
 
 __all__ = [
     "RandomAgentsMultiPolicy",
     "ThinkyAgentsMultiPolicy",
     "RaceCarAgentsMultiPolicy",
     "LadyBugAgentsMultiPolicy",
+    "CogsguardAlignAllAgentsMultiPolicy",
 ]
-
-# Re-export the policy classes for convenience
-from cogames.policy.nim_agents.agents import (  # noqa: F401
-    LadyBugAgentsMultiPolicy,
-    RaceCarAgentsMultiPolicy,
-    RandomAgentsMultiPolicy,
-    ThinkyAgentsMultiPolicy,
-)

cogames/policy/nim_agents/agents.py

@@ -1,66 +1,127 @@
-import importlib
-import os
-import sys
-from typing import Sequence
-
-from mettagrid.policy.policy import NimMultiAgentPolicy
-from mettagrid.policy.policy_env_interface import PolicyEnvInterface
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-bindings_dir = os.path.join(current_dir, "bindings/generated")
-if bindings_dir not in sys.path:
-    sys.path.append(bindings_dir)
-
-na = importlib.import_module("nim_agents")
-
-
-def start_measure():
-    na.start_measure()
-
-
-def end_measure():
-    na.end_measure()
+"""Legacy module path for Nim-based agent policies.
 
+The implementation moved to `cogames_agents.policy.nim_agents.agents`. Keep this
+wrapper so older checkpoints and policy bundles that reference the old import
+path can still load.
+"""
 
-class ThinkyAgentsMultiPolicy(NimMultiAgentPolicy):
-    short_names = ["thinky"]
+from __future__ import annotations
 
-    def __init__(self, policy_env_info: PolicyEnvInterface, agent_ids: Sequence[int] | None = None):
-        super().__init__(
-            policy_env_info,
-            nim_policy_factory=na.ThinkyPolicy,
-            agent_ids=agent_ids,
-        )
-
-
-class RandomAgentsMultiPolicy(NimMultiAgentPolicy):
-    short_names = ["nim_random"]
-
-    def __init__(self, policy_env_info: PolicyEnvInterface, agent_ids: Sequence[int] | None = None):
-        super().__init__(
-            policy_env_info,
-            nim_policy_factory=na.RandomPolicy,
-            agent_ids=agent_ids,
-        )
-
-
-class RaceCarAgentsMultiPolicy(NimMultiAgentPolicy):
-    short_names = ["race_car"]
-
-    def __init__(self, policy_env_info: PolicyEnvInterface, agent_ids: Sequence[int] | None = None):
-        super().__init__(
-            policy_env_info,
-            nim_policy_factory=na.RaceCarPolicy,
-            agent_ids=agent_ids,
-        )
+from collections.abc import Callable
 
+from mettagrid.policy.policy import MultiAgentPolicy
+from mettagrid.policy.policy_env_interface import PolicyEnvInterface
 
-class LadyBugAgentsMultiPolicy(NimMultiAgentPolicy):
-    short_names = ["nim_ladybug"]
 
-    def __init__(self, policy_env_info: PolicyEnvInterface, agent_ids: Sequence[int] | None = None):
-        super().__init__(
-            policy_env_info,
-            nim_policy_factory=na.LadybugPolicy,
-            agent_ids=agent_ids,
-        )
+def _raise_missing_nim_agents() -> None:
+    raise ModuleNotFoundError(
+        "Nim scripted agents are not available (missing `nim_agents` bindings). "
+        "If you're developing locally, build them with: "
+        "`cd packages/cogames-agents/src/cogames_agents/policy/nim_agents && nim c nim_agents.nim`."
+    )
+
+
+start_measure: Callable[[], None]
+end_measure: Callable[[], None]
+ThinkyAgentsMultiPolicy: type[MultiAgentPolicy]
+RandomAgentsMultiPolicy: type[MultiAgentPolicy]
+RaceCarAgentsMultiPolicy: type[MultiAgentPolicy]
+LadyBugAgentsMultiPolicy: type[MultiAgentPolicy]
+CogsguardAgentsMultiPolicy: type[MultiAgentPolicy]
+CogsguardAlignAllAgentsMultiPolicy: type[MultiAgentPolicy]
+
+try:
+    from cogames_agents.policy.nim_agents.agents import (
+        CogsguardAgentsMultiPolicy as _CogsguardAgentsMultiPolicy,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        CogsguardAlignAllAgentsMultiPolicy as _CogsguardAlignAllAgentsMultiPolicy,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        LadyBugAgentsMultiPolicy as _LadyBugAgentsMultiPolicy,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        RaceCarAgentsMultiPolicy as _RaceCarAgentsMultiPolicy,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        RandomAgentsMultiPolicy as _RandomAgentsMultiPolicy,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        ThinkyAgentsMultiPolicy as _ThinkyAgentsMultiPolicy,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        end_measure as _end_measure,
+    )
+    from cogames_agents.policy.nim_agents.agents import (
+        start_measure as _start_measure,
+    )
+except (ModuleNotFoundError, OSError) as exc:
+    if (
+        isinstance(exc, ModuleNotFoundError)
+        and exc.name
+        and (exc.name == "cogames_agents" or exc.name.startswith("cogames_agents."))
+    ):
+        raise ModuleNotFoundError(
+            "Legacy import `cogames.policy.nim_agents.agents` requires optional dependency "
+            "`cogames-agents` (install `cogames[agents]`)."
+        ) from exc
+    # Fall back to stubs if the optional Nim bindings are missing.
+
+    def _missing_start_measure() -> None:
+        _raise_missing_nim_agents()
+
+    def _missing_end_measure() -> None:
+        _raise_missing_nim_agents()
+
+    class _MissingThinkyAgentsMultiPolicy(MultiAgentPolicy):
+        def __init__(self, policy_env_info: PolicyEnvInterface, **_: object):
+            _raise_missing_nim_agents()
+
+    class _MissingRandomAgentsMultiPolicy(MultiAgentPolicy):
+        def __init__(self, policy_env_info: PolicyEnvInterface, **_: object):
+            _raise_missing_nim_agents()
+
+    class _MissingRaceCarAgentsMultiPolicy(MultiAgentPolicy):
+        def __init__(self, policy_env_info: PolicyEnvInterface, **_: object):
+            _raise_missing_nim_agents()
+
+    class _MissingLadyBugAgentsMultiPolicy(MultiAgentPolicy):
+        def __init__(self, policy_env_info: PolicyEnvInterface, **_: object):
+            _raise_missing_nim_agents()
+
+    class _MissingCogsguardAgentsMultiPolicy(MultiAgentPolicy):
+        def __init__(self, policy_env_info: PolicyEnvInterface, **_: object):
+            _raise_missing_nim_agents()
+
+    class _MissingCogsguardAlignAllAgentsMultiPolicy(MultiAgentPolicy):
+        def __init__(self, policy_env_info: PolicyEnvInterface, **_: object):
+            _raise_missing_nim_agents()
+
+    start_measure = _missing_start_measure
+    end_measure = _missing_end_measure
+    ThinkyAgentsMultiPolicy = _MissingThinkyAgentsMultiPolicy
+    RandomAgentsMultiPolicy = _MissingRandomAgentsMultiPolicy
+    RaceCarAgentsMultiPolicy = _MissingRaceCarAgentsMultiPolicy
+    LadyBugAgentsMultiPolicy = _MissingLadyBugAgentsMultiPolicy
+    CogsguardAgentsMultiPolicy = _MissingCogsguardAgentsMultiPolicy
+    CogsguardAlignAllAgentsMultiPolicy = _MissingCogsguardAlignAllAgentsMultiPolicy
+else:
+    start_measure = _start_measure
+    end_measure = _end_measure
+    ThinkyAgentsMultiPolicy = _ThinkyAgentsMultiPolicy
+    RandomAgentsMultiPolicy = _RandomAgentsMultiPolicy
+    RaceCarAgentsMultiPolicy = _RaceCarAgentsMultiPolicy
+    LadyBugAgentsMultiPolicy = _LadyBugAgentsMultiPolicy
+    CogsguardAgentsMultiPolicy = _CogsguardAgentsMultiPolicy
+    CogsguardAlignAllAgentsMultiPolicy = _CogsguardAlignAllAgentsMultiPolicy
+
+__all__ = [
+    "start_measure",
+    "end_measure",
+    "ThinkyAgentsMultiPolicy",
+    "RandomAgentsMultiPolicy",
+    "RaceCarAgentsMultiPolicy",
+    "LadyBugAgentsMultiPolicy",
+    "CogsguardAgentsMultiPolicy",
+    "CogsguardAlignAllAgentsMultiPolicy",
+]