cogames 0.3.65__py3-none-any.whl → 0.3.69__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,10 +8,8 @@ from rich import box
 from rich.table import Table
 
 from cogames.cli.base import console
-from cogames.cogs_vs_clips.mission import (
-    CvCMission,
-    NumCogsVariant,
-)
+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
@@ -27,16 +25,16 @@ from mettagrid.mapgen.mapgen import MapGen
 
 @lru_cache(maxsize=1)
 def _get_core_missions() -> list[CvCMission]:
-    from cogames.cogs_vs_clips.missions import get_core_missions
+    from cogames.cogs_vs_clips.missions import get_core_missions  # noqa: PLC0415
 
     return get_core_missions()
 
 
 @lru_cache(maxsize=1)
 def _get_eval_missions_all() -> list[CvCMission]:
-    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
+    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
 
     missions: list[CvCMission] = []
     missions.extend(INTEGRATED_EVAL_MISSIONS)
@@ -74,19 +72,21 @@ def load_mission_set(mission_set: str) -> list[CvCMission]:
                 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:
@@ -132,6 +132,12 @@ def parse_variants(variants_arg: Optional[list[str]]) -> list[CoGameMissionVaria
     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()]
@@ -157,17 +163,25 @@ def get_site_by_name(site_name: str) -> CoGameSite:
 
 
 def get_mission_name_and_config(
-    ctx: typer.Context, mission_arg: Optional[str], variants_arg: Optional[list[str]] = None, cogs: Optional[int] = None
+    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()
 
@@ -183,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())
@@ -192,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 = []
@@ -224,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
@@ -232,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)]
 
 
@@ -278,6 +296,7 @@ def get_mission(
     variants_arg: Optional[list[str]] = None,
     cogs: Optional[int] = None,
     include_legacy: bool = False,
+    difficulty: Optional[str] = None,
 ) -> tuple[str, MettaGridConfig, Optional[CvCMission]]:
     """Get a specific mission configuration by name or file path.
 
@@ -286,6 +305,7 @@ 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, CvCMission or None)
@@ -295,6 +315,8 @@ def get_mission(
     """
     # 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}")
@@ -311,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:
@@ -322,6 +345,8 @@ def get_mission(
 
     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:

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)