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

cogames/cli/client.py

@@ -253,9 +253,6 @@ class TournamentServerClient:
             json=payload,
         )
 
-    def update_policy_version_tags(self, policy_version_id: uuid.UUID, tags: dict[str, str]) -> dict[str, Any]:
-        return self._put(f"/stats/policies/versions/{policy_version_id}/tags", json=tags)
-
     def get_policy_memberships(self, policy_version_id: uuid.UUID) -> list[dict[str, Any]]:
         return self._get(f"/tournament/policies/{policy_version_id}/memberships", list[dict[str, Any]])
 

cogames/cli/docsync/docsync.py

@@ -120,7 +120,13 @@ def check(
         typer.echo("\nNotebook documentation is out of sync!", err=True)
         for error in errors:
             typer.echo(f"  {error}", err=True)
-        typer.echo("\nTo fix: run 'cogames docsync all'", err=True)
+        typer.echo(
+            "\nThis can happen if you modified cogames CLI flags (the README includes a command reference)"
+            "\nor changed code that affects notebook outputs."
+            "\n"
+            "\nTo fix: run 'cogames docsync all'",
+            err=True,
+        )
         raise typer.Exit(1)
 
     typer.echo("All notebooks are in sync!")

cogames/cli/mission.py

@@ -8,52 +8,42 @@ from rich import box
 from rich.table import Table
 
 from cogames.cli.base import console
-from cogames.cogs_vs_clips.mission import (
-    MAP_MISSION_DELIMITER,
-    AnyMission,
-    Mission,
-    MissionVariant,
-    NumCogsVariant,
-    Site,
-)
-from cogames.cogs_vs_clips.procedural import MachinaArena
+from cogames.cogs_vs_clips.clip_difficulty import get_cogsguard_difficulty
+from cogames.cogs_vs_clips.mission import CvCMission, NumCogsVariant
 from cogames.cogs_vs_clips.sites import SITES
+from cogames.cogs_vs_clips.terrain import MachinaArena
 from cogames.cogs_vs_clips.variants import HIDDEN_VARIANTS, VARIANTS
+from cogames.core import (
+    MAP_MISSION_DELIMITER,
+    CoGameMissionVariant,
+    CoGameSite,
+)
 from cogames.game import load_mission_config, load_mission_config_from_python
 from mettagrid import MettaGridConfig
 from mettagrid.mapgen.mapgen import MapGen
 
 
 @lru_cache(maxsize=1)
-def _get_core_missions() -> list[AnyMission]:
-    from cogames.cogs_vs_clips.missions import get_core_missions
+def _get_core_missions() -> list[CvCMission]:
+    from cogames.cogs_vs_clips.missions import get_core_missions  # noqa: PLC0415
 
     return get_core_missions()
 
 
 @lru_cache(maxsize=1)
-def _get_legacy_missions() -> list[Mission]:
-    from cogames.cogs_vs_clips.missions import get_legacy_missions
+def _get_eval_missions_all() -> list[CvCMission]:
+    from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS  # noqa: PLC0415
+    from cogames.cogs_vs_clips.evals.integrated_evals import EVAL_MISSIONS as INTEGRATED_EVAL_MISSIONS  # noqa: PLC0415
+    from cogames.cogs_vs_clips.evals.spanning_evals import EVAL_MISSIONS as SPANNING_EVAL_MISSIONS  # noqa: PLC0415
 
-    return get_legacy_missions()
-
-
-@lru_cache(maxsize=1)
-def _get_eval_missions_all() -> list[AnyMission]:
-    from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS
-    from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS
-    from cogames.cogs_vs_clips.evals.integrated_evals import EVAL_MISSIONS as INTEGRATED_EVAL_MISSIONS
-    from cogames.cogs_vs_clips.evals.spanning_evals import EVAL_MISSIONS as SPANNING_EVAL_MISSIONS
-
-    missions: list[AnyMission] = []
-    missions.extend(COGSGUARD_EVAL_MISSIONS)
+    missions: list[CvCMission] = []
     missions.extend(INTEGRATED_EVAL_MISSIONS)
     missions.extend(SPANNING_EVAL_MISSIONS)
     missions.extend(mission_cls() for mission_cls in DIAGNOSTIC_EVALS)  # type: ignore[call-arg]
     return missions
 
 
-def load_mission_set(mission_set: str) -> list[AnyMission]:
+def load_mission_set(mission_set: str) -> list[CvCMission]:
     """Load a predefined set of evaluation missions.
 
     Args:
@@ -70,7 +60,7 @@ def load_mission_set(mission_set: str) -> list[AnyMission]:
     Raises:
         ValueError: If mission_set name is unknown
     """
-    missions_list: list[AnyMission]
+    missions_list: list[CvCMission]
     if mission_set == "all":
         # All missions: eval missions + integrated + spanning + diagnostic + core missions
         missions_list = list(_get_eval_missions_all())
@@ -82,19 +72,21 @@ def load_mission_set(mission_set: str) -> list[AnyMission]:
                 missions_list.append(mission)
 
     elif mission_set == "diagnostic_evals":
-        from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS
+        from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS  # noqa: PLC0415
 
         missions_list = [mission_cls() for mission_cls in DIAGNOSTIC_EVALS]  # type: ignore[call-arg]
     elif mission_set == "cogsguard_evals":
-        from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS
+        from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS  # noqa: PLC0415
 
         missions_list = list(COGSGUARD_EVAL_MISSIONS)
     elif mission_set == "integrated_evals":
-        from cogames.cogs_vs_clips.evals.integrated_evals import EVAL_MISSIONS as INTEGRATED_EVAL_MISSIONS
+        from cogames.cogs_vs_clips.evals.integrated_evals import (  # noqa: PLC0415
+            EVAL_MISSIONS as INTEGRATED_EVAL_MISSIONS,
+        )
 
         missions_list = list(INTEGRATED_EVAL_MISSIONS)
     elif mission_set == "spanning_evals":
-        from cogames.cogs_vs_clips.evals.spanning_evals import EVAL_MISSIONS as SPANNING_EVAL_MISSIONS
+        from cogames.cogs_vs_clips.evals.spanning_evals import EVAL_MISSIONS as SPANNING_EVAL_MISSIONS  # noqa: PLC0415
 
         missions_list = list(SPANNING_EVAL_MISSIONS)
     else:
@@ -104,14 +96,14 @@ def load_mission_set(mission_set: str) -> list[AnyMission]:
     return missions_list
 
 
-def parse_variants(variants_arg: Optional[list[str]]) -> list[MissionVariant]:
+def parse_variants(variants_arg: Optional[list[str]]) -> list[CoGameMissionVariant]:
     """Parse variant specifications from command line.
 
     Args:
         variants_arg: List of variant names like ["solar_flare", "dark_side"]
 
     Returns:
-        List of configured MissionVariant instances
+        List of configured CoGameMissionVariant instances
 
     Raises:
         ValueError: If variant name is unknown
@@ -119,11 +111,11 @@ def parse_variants(variants_arg: Optional[list[str]]) -> list[MissionVariant]:
     if not variants_arg:
         return []
 
-    variants: list[MissionVariant] = []
+    variants: list[CoGameMissionVariant] = []
     all_variants = [*VARIANTS, *HIDDEN_VARIANTS]
     for name in variants_arg:
         # Find matching variant class by instantiating and checking the name
-        variant: MissionVariant | None = None
+        variant: CoGameMissionVariant | None = None
         for v in all_variants:
             if v.name == name:
                 variant = v
@@ -140,6 +132,12 @@ def parse_variants(variants_arg: Optional[list[str]]) -> list[MissionVariant]:
     return variants
 
 
+def parse_difficulty(difficulty_arg: Optional[str]) -> Optional[CoGameMissionVariant]:
+    if difficulty_arg is None:
+        return None
+    return get_cogsguard_difficulty(difficulty_arg)
+
+
 def get_all_missions() -> list[str]:
     """Get all core mission names in the format site.mission (excludes evals)."""
     return [mission.full_name() for mission in _get_core_missions()]
@@ -150,7 +148,7 @@ def get_all_eval_missions() -> list[str]:
     return [mission.full_name() for mission in _get_eval_missions_all()]
 
 
-def get_site_by_name(site_name: str) -> Site:
+def get_site_by_name(site_name: str) -> CoGameSite:
     """Get a site by name.
 
     Raises:
@@ -165,17 +163,25 @@ def get_site_by_name(site_name: str) -> Site:
 
 
 def get_mission_name_and_config(
-    ctx: typer.Context, mission_arg: Optional[str], variants_arg: Optional[list[str]] = None, cogs: Optional[int] = None
-) -> tuple[str, MettaGridConfig, Optional[AnyMission]]:
+    ctx: typer.Context,
+    mission_arg: Optional[str],
+    variants_arg: Optional[list[str]] = None,
+    cogs: Optional[int] = None,
+    difficulty: Optional[str] = None,
+) -> tuple[str, MettaGridConfig, Optional[CvCMission]]:
     if not mission_arg:
         console.print(ctx.get_help())
         console.print("[yellow]Missing: --mission / -m[/yellow]\n")
     else:
         try:
-            return get_mission(mission_arg, variants_arg, cogs)
+            return get_mission(mission_arg, variants_arg=variants_arg, cogs=cogs, difficulty=difficulty)
         except ValueError as e:
-            console.print(f"[red]Mission '{mission_arg}' not found.[/red]")
-            console.print("[dim]Run: cogames missions (or cogames missions <site>) to list options.[/dim]\n")
+            error_msg = str(e)
+            if "variant" in error_msg.lower() or "difficulty" in error_msg.lower():
+                console.print(f"[red]{error_msg}[/red]")
+            else:
+                console.print(f"[red]Mission '{mission_arg}' not found.[/red]")
+                console.print("[dim]Run: cogames missions (or cogames missions <site>) to list options.[/dim]\n")
             raise typer.Exit(1) from e
     list_missions()
 
@@ -191,6 +197,7 @@ def get_mission_names_and_configs(
     variants_arg: Optional[list[str]] = None,
     cogs: Optional[int] = None,
     steps: Optional[int] = None,
+    difficulty: Optional[str] = None,
 ) -> list[tuple[str, MettaGridConfig]]:
     if not missions_arg:
         console.print(ctx.get_help())
@@ -200,7 +207,7 @@ def get_mission_names_and_configs(
             not_deduped = [
                 mission
                 for missions in missions_arg
-                for mission in _get_missions_by_possible_wildcard(missions, variants_arg, cogs)
+                for mission in _get_missions_by_possible_wildcard(missions, variants_arg, cogs, difficulty)
             ]
             name_set: set[str] = set()
             deduped = []
@@ -232,6 +239,7 @@ def _get_missions_by_possible_wildcard(
     mission_arg: str,
     variants_arg: Optional[list[str]],
     cogs: Optional[int],
+    difficulty: Optional[str],
 ) -> list[tuple[str, MettaGridConfig]]:
     if "*" in mission_arg:
         # Convert shell-style wildcard to regex pattern
@@ -240,10 +248,12 @@ def _get_missions_by_possible_wildcard(
         # Drop the Mission (3rd element) for wildcard results
         return [
             (name, env_cfg)
-            for name, env_cfg, _ in (get_mission(m, variants_arg=variants_arg, cogs=cogs) for m in missions)
+            for name, env_cfg, _ in (
+                get_mission(m, variants_arg=variants_arg, cogs=cogs, difficulty=difficulty) for m in missions
+            )
         ]
     # Drop the Mission for single mission
-    name, env_cfg, _ = get_mission(mission_arg, variants_arg=variants_arg, cogs=cogs)
+    name, env_cfg, _ = get_mission(mission_arg, variants_arg=variants_arg, cogs=cogs, difficulty=difficulty)
     return [(name, env_cfg)]
 
 
@@ -253,12 +263,10 @@ def find_mission(
     *,
     include_evals: bool = False,
     include_legacy: bool = False,
-) -> AnyMission:
-    missions: list[AnyMission] = list(_get_core_missions())
+) -> CvCMission:
+    missions: list[CvCMission] = list(_get_core_missions())
     if include_evals:
         missions = [*missions, *_get_eval_missions_all()]
-    if include_legacy:
-        missions = [*missions, *_get_legacy_missions()]
 
     found_site = False
     for mission in missions:
@@ -288,7 +296,8 @@ def get_mission(
     variants_arg: Optional[list[str]] = None,
     cogs: Optional[int] = None,
     include_legacy: bool = False,
-) -> tuple[str, MettaGridConfig, Optional[AnyMission]]:
+    difficulty: Optional[str] = None,
+) -> tuple[str, MettaGridConfig, Optional[CvCMission]]:
     """Get a specific mission configuration by name or file path.
 
     Args:
@@ -296,15 +305,18 @@ def get_mission(
         variants_arg: List of variant names like ["solar_flare", "dark_side"]
         cogs: Number of cogs (agents) to use, overrides the default from the mission
         include_legacy: Whether to include legacy (pre-CogsGuard) missions
+        difficulty: Difficulty name (easy, medium, hard) controlling clips events
 
     Returns:
-        Tuple of (mission name, MettaGridConfig, Mission or None)
+        Tuple of (mission name, MettaGridConfig, CvCMission or None)
 
     Raises:
         ValueError: If mission not found or file cannot be loaded
     """
     # Check if it's a file path
     if any(mission_arg.endswith(ext) for ext in [".yaml", ".yml", ".json", ".py"]):
+        if difficulty is not None:
+            raise ValueError("Difficulty can only be used with mission names, not config files.")
         path = Path(mission_arg)
         if not path.exists():
             raise ValueError(f"File not found: {mission_arg}")
@@ -321,6 +333,7 @@ def get_mission(
 
     # Parse variants if provided
     variants = parse_variants(variants_arg)
+    difficulty_variant = parse_difficulty(difficulty)
 
     # Otherwise, treat it as a fully qualified mission name, or as a site name
     if (delim_count := mission_arg.count(MAP_MISSION_DELIMITER)) == 0:
@@ -330,8 +343,10 @@ def get_mission(
     else:
         site_name, mission_name = mission_arg.split(MAP_MISSION_DELIMITER)
 
-    mission: AnyMission = find_mission(site_name, mission_name, include_evals=True, include_legacy=include_legacy)
+    mission: CvCMission = find_mission(site_name, mission_name, include_evals=True, include_legacy=include_legacy)
 
+    if difficulty_variant is not None:
+        mission = mission.with_variants([difficulty_variant])
     if variants:
         mission = mission.with_variants(variants)
     if cogs is not None:
@@ -454,7 +469,7 @@ def list_evals() -> None:
         return
 
     # Group missions by site
-    missions_by_site: dict[str, list[AnyMission]] = {}
+    missions_by_site: dict[str, list[CvCMission]] = {}
     for m in evals:
         missions_by_site.setdefault(m.site.name, []).append(m)
 
@@ -506,7 +521,7 @@ def list_evals() -> None:
     console.print("  [bold]cogames play[/bold] --mission [blue]evals.divide_and_conquer[/blue]")
 
 
-def describe_mission(mission_name: str, game_config: MettaGridConfig, mission_cfg: AnyMission | None = None) -> None:
+def describe_mission(mission_name: str, game_config: MettaGridConfig, mission_cfg: CvCMission | None = None) -> None:
     """Print detailed information about a specific mission.
 
     Args:

cogames/cli/policy.py

@@ -68,13 +68,13 @@ def _translate_error(e: Exception) -> str:
     return translated
 
 
-def get_policy_spec(ctx: typer.Context, policy_arg: Optional[str]) -> PolicySpec:
+def get_policy_spec(ctx: typer.Context, policy_arg: Optional[str], device: str | None = None) -> PolicySpec:
     if policy_arg is None:
         console.print(ctx.get_help())
         console.print("[yellow]Missing: --policy / -p[/yellow]\n")
     else:
         try:
-            return parse_policy_spec(spec=policy_arg).to_policy_spec()
+            return parse_policy_spec(spec=policy_arg, device=device).to_policy_spec()
         except (ValueError, ModuleNotFoundError) as e:
             translated = _translate_error(e)
             console.print(f"[yellow]Error parsing policy argument: {translated}[/yellow]\n")
@@ -90,14 +90,14 @@ def get_policy_spec(ctx: typer.Context, policy_arg: Optional[str]) -> PolicySpec
 
 
 def get_policy_specs_with_proportions(
-    ctx: typer.Context, policy_args: Optional[list[str]]
+    ctx: typer.Context, policy_args: Optional[list[str]], device: str | None = None
 ) -> list[PolicySpecWithProportion]:
     if not policy_args:
         console.print(ctx.get_help())
         console.print("[yellow]Supply at least one: --policy / -p[/yellow]\n")
     else:
         try:
-            return [parse_policy_spec(spec=policy_arg) for policy_arg in policy_args]
+            return [parse_policy_spec(spec=policy_arg, device=device) for policy_arg in policy_args]
         except (ValueError, ModuleNotFoundError) as e:
             translated = _translate_error(e)
             console.print(f"[yellow]Error parsing policy argument: {translated}[/yellow]")
@@ -112,7 +112,20 @@ def get_policy_specs_with_proportions(
     raise typer.Exit(1)
 
 
-def parse_policy_spec(spec: str) -> PolicySpecWithProportion:
+def _apply_device_override(spec: PolicySpecWithProportion, device: str | None) -> PolicySpecWithProportion:
+    if device is None:
+        return spec
+    init_kwargs = dict(spec.init_kwargs or {})
+    init_kwargs["device"] = device
+    return PolicySpecWithProportion(
+        class_path=spec.class_path,
+        data_path=spec.data_path,
+        proportion=spec.proportion,
+        init_kwargs=init_kwargs,
+    )
+
+
+def parse_policy_spec(spec: str, device: str | None = None) -> PolicySpecWithProportion:
     """Parse a policy CLI option into its components.
 
     Supports two formats:
@@ -161,13 +174,14 @@ def parse_policy_spec(spec: str) -> PolicySpecWithProportion:
             uri_part = spec
             fraction = 1.0
 
-        policy = policy_spec_from_uri(uri_part.strip())
-        return PolicySpecWithProportion(
+        policy = policy_spec_from_uri(uri_part.strip(), device=device or "cpu")
+        policy_spec = PolicySpecWithProportion(
             class_path=policy.class_path,
             data_path=policy.data_path,
             proportion=fraction,
             init_kwargs=policy.init_kwargs,
         )
+        return _apply_device_override(policy_spec, device)
 
     entries = [part.strip() for part in spec.split(",") if part.strip()]
     if not entries:
@@ -176,19 +190,20 @@ def parse_policy_spec(spec: str) -> PolicySpecWithProportion:
     fraction = 1.0
     first = entries[0]
     if is_uri(first) or ("=" not in first and is_path_like(first)):
-        policy = policy_spec_from_uri(first)
+        policy = policy_spec_from_uri(first, device=device or "cpu")
         for entry in entries[1:]:
             key, value = parse_key_value(entry)
             if key != "proportion":
                 raise ValueError("Only proportion is supported after a checkpoint URI.")
             fraction = parse_proportion(value)
 
-        return PolicySpecWithProportion(
+        policy_spec = PolicySpecWithProportion(
             class_path=policy.class_path,
             data_path=policy.data_path,
             proportion=fraction,
             init_kwargs=policy.init_kwargs,
         )
+        return _apply_device_override(policy_spec, device)
 
     if "=" not in first:
         if ":" in first:
@@ -230,9 +245,10 @@ def parse_policy_spec(spec: str) -> PolicySpecWithProportion:
     if class_path is None:
         raise ValueError("Policy specification must include class= for key=value format.")
 
-    return PolicySpecWithProportion(
+    policy_spec = PolicySpecWithProportion(
         class_path=class_path,
         data_path=data_path,
         proportion=fraction,
         init_kwargs=init_kwargs,
     )
+    return _apply_device_override(policy_spec, device)