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

cogames/main.py

@@ -8,6 +8,7 @@ from cogames.cli.utils import suppress_noisy_logs
 
 suppress_noisy_logs()
 
+import importlib
 import importlib.metadata
 import importlib.util
 import json
@@ -17,6 +18,7 @@ import subprocess
 import sys
 import threading
 import time
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Literal, Optional, TypeVar
 
@@ -38,7 +40,6 @@ from cogames import play as play_module
 from cogames import train as train_module
 from cogames.cli.base import console
 from cogames.cli.client import SeasonInfo, TournamentServerClient, fetch_default_season, fetch_season_info
-from cogames.cli.docsync import docsync
 from cogames.cli.leaderboard import (
     leaderboard_cmd,
     parse_policy_identifier,
@@ -63,6 +64,7 @@ from cogames.cli.policy import (
     policy_arg_w_proportion_example,
 )
 from cogames.cli.submit import DEFAULT_SUBMIT_SERVER, results_url_for_season, upload_policy, validate_policy_spec
+from cogames.cogs_vs_clips.mission import CvCMission, NumCogsVariant
 from cogames.curricula import make_rotation
 from cogames.device import resolve_training_device
 from mettagrid.config.mettagrid_config import MettaGridConfig
@@ -87,6 +89,158 @@ logger = logging.getLogger("cogames.main")
 T = TypeVar("T")
 
 
+@dataclass(frozen=True)
+class DiagnoseCase:
+    name: str
+    env_cfg: MettaGridConfig
+
+
+def _load_eval_missions(module_path: str) -> list[CvCMission]:
+    module = importlib.import_module(module_path)
+    missions = getattr(module, "EVAL_MISSIONS", None)
+    if missions is None:
+        raise AttributeError(f"Module '{module_path}' does not define EVAL_MISSIONS")
+    return list(missions)
+
+
+def _load_diagnose_missions(mission_set: str) -> list[CvCMission]:
+    if mission_set == "thinky_evals":
+        return []
+
+    if mission_set == "all":
+        from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS  # noqa: PLC0415
+        from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS  # noqa: PLC0415
+        from cogames.cogs_vs_clips.missions import MISSIONS as ALL_MISSIONS  # noqa: PLC0415
+
+        missions_list: list[CvCMission] = []
+        missions_list.extend(COGSGUARD_EVAL_MISSIONS)
+        missions_list.extend(_load_eval_missions("cogames.cogs_vs_clips.evals.integrated_evals"))
+        missions_list.extend(_load_eval_missions("cogames.cogs_vs_clips.evals.spanning_evals"))
+        missions_list.extend([mission_cls() for mission_cls in DIAGNOSTIC_EVALS])  # type: ignore[call-arg]
+        eval_mission_names = {mission.name for mission in missions_list}
+        for mission in ALL_MISSIONS:
+            if mission.name not in eval_mission_names:
+                missions_list.append(mission)
+        return missions_list
+
+    if mission_set == "cogsguard_evals":
+        from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS  # noqa: PLC0415
+
+        return list(COGSGUARD_EVAL_MISSIONS)
+
+    if mission_set == "diagnostic_evals":
+        from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS  # noqa: PLC0415
+
+        return [mission_cls() for mission_cls in DIAGNOSTIC_EVALS]  # type: ignore[call-arg]
+
+    if mission_set == "tournament":
+        from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS  # noqa: PLC0415
+
+        missions_list = []
+        missions_list.extend(_load_eval_missions("cogames.cogs_vs_clips.evals.integrated_evals"))
+        missions_list.extend([mission_cls() for mission_cls in DIAGNOSTIC_EVALS])  # type: ignore[call-arg]
+        return missions_list
+
+    if mission_set == "integrated_evals":
+        return _load_eval_missions("cogames.cogs_vs_clips.evals.integrated_evals")
+
+    if mission_set == "spanning_evals":
+        return _load_eval_missions("cogames.cogs_vs_clips.evals.spanning_evals")
+
+    raise ValueError(f"Unknown mission set: {mission_set}")
+
+
+def _build_thinky_mission_map() -> dict[str, CvCMission]:
+    from cogames.cogs_vs_clips.evals.cogsguard_evals import COGSGUARD_EVAL_MISSIONS  # noqa: PLC0415
+    from cogames.cogs_vs_clips.evals.diagnostic_evals import DIAGNOSTIC_EVALS  # noqa: PLC0415
+    from cogames.cogs_vs_clips.missions import MISSIONS as ALL_MISSIONS  # noqa: PLC0415
+
+    missions: list[CvCMission] = []
+    missions.extend(_load_eval_missions("cogames.cogs_vs_clips.evals.integrated_evals"))
+    missions.extend(_load_eval_missions("cogames.cogs_vs_clips.evals.spanning_evals"))
+    missions.extend([mission_cls() for mission_cls in DIAGNOSTIC_EVALS])  # type: ignore[call-arg]
+    missions.extend(COGSGUARD_EVAL_MISSIONS)
+    missions.extend(ALL_MISSIONS)
+
+    mission_map: dict[str, CvCMission] = {}
+    for mission in missions:
+        mission_map.setdefault(mission.name, mission)
+    return mission_map
+
+
+def _matches_experiment(mission_name: str, experiment_filters: set[str]) -> bool:
+    if not experiment_filters:
+        return True
+    if mission_name in experiment_filters:
+        return True
+    suffix = f".{mission_name}"
+    return any(name.endswith(suffix) for name in experiment_filters)
+
+
+def _cogs_for_mission(mission: CvCMission, cogs_list: list[int], respect_cogs_list: bool) -> list[int]:
+    fixed_cogs = getattr(mission, "num_cogs", None)
+    if fixed_cogs is not None:
+        if respect_cogs_list and fixed_cogs not in cogs_list:
+            return []
+        return [fixed_cogs]
+    site = getattr(mission, "site", None)
+    if site is None:
+        return list(cogs_list)
+    min_cogs = getattr(site, "min_cogs", None)
+    max_cogs = getattr(site, "max_cogs", None)
+    return [
+        num_cogs
+        for num_cogs in cogs_list
+        if (min_cogs is None or num_cogs >= min_cogs) and (max_cogs is None or num_cogs <= max_cogs)
+    ]
+
+
+def _build_diagnose_case(mission: CvCMission, num_cogs: int, steps: int) -> DiagnoseCase:
+    mission_with_cogs = mission.with_variants([NumCogsVariant(num_cogs=num_cogs)])
+    env_cfg = mission_with_cogs.make_env()
+    env_cfg.game.max_steps = steps
+    name = f"{mission.full_name()} (cogs={num_cogs})"
+    return DiagnoseCase(name=name, env_cfg=env_cfg)
+
+
+def _build_diagnose_cases(
+    *,
+    mission_set: str,
+    experiments: Optional[list[str]],
+    cogs: Optional[list[int]],
+    steps: int,
+) -> list[DiagnoseCase]:
+    experiment_filters = set(experiments or [])
+    cogs_list = cogs if cogs else [1, 2, 4]
+    respect_cogs_list = cogs is not None
+    cases: list[DiagnoseCase] = []
+
+    if mission_set == "thinky_evals":
+        from cogames_agents.policy.nim_agents.thinky_eval import EVALS as THINKY_EVALS  # noqa: PLC0415
+
+        mission_map = _build_thinky_mission_map()
+        for exp_name, _tag, num_cogs in THINKY_EVALS:
+            if not _matches_experiment(exp_name, experiment_filters):
+                continue
+            if respect_cogs_list and num_cogs not in cogs_list:
+                continue
+            base_mission = mission_map.get(exp_name)
+            if base_mission is None:
+                logger.warning("Thinky eval mission '%s' not found; skipping.", exp_name)
+                continue
+            cases.append(_build_diagnose_case(base_mission, num_cogs, steps))
+        return cases
+
+    missions = _load_diagnose_missions(mission_set)
+    for mission in missions:
+        if not _matches_experiment(mission.name, experiment_filters):
+            continue
+        for num_cogs in _cogs_for_mission(mission, cogs_list, respect_cogs_list):
+            cases.append(_build_diagnose_case(mission, num_cogs, steps))
+
+    return cases
+
+
 def _resolve_mettascope_script() -> Path:
     spec = importlib.util.find_spec("mettagrid")
     if spec is None or spec.origin is None:
@@ -132,7 +286,18 @@ tutorial_app = typer.Typer(
 if register_tribal_cli is not None:
     register_tribal_cli(app)
 
-app.add_typer(docsync.app, name="docsync", hidden=True)
+
+@app.command(
+    name="docsync",
+    hidden=True,
+    context_settings={"allow_extra_args": True, "ignore_unknown_options": True},
+    add_help_option=False,
+)
+def docsync_cmd(ctx: typer.Context) -> None:
+    """Sync cogames docs between .ipynb, .py, and .md formats (dev-only)."""
+    from cogames.cli.docsync import docsync  # noqa: PLC0415
+
+    docsync.app(prog_name="cogames docsync", standalone_mode=False, args=list(ctx.args))
 
 
 @tutorial_app.command(
@@ -160,7 +325,7 @@ def tutorial_cmd(
     console.print("[dim]Initializing Mettascope...[/dim]")
 
     # Load tutorial mission (CogsGuard)
-    from cogames.cogs_vs_clips.missions import make_cogsguard_mission
+    from cogames.cogs_vs_clips.missions import make_cogsguard_mission  # noqa: PLC0415
 
     # Create environment config
     env_cfg = make_cogsguard_mission(num_agents=1, max_steps=1000).make_env()
@@ -310,7 +475,7 @@ def cogsguard_tutorial_cmd(
     console.print("[dim]Initializing Mettascope...[/dim]")
 
     # Load CogsGuard tutorial mission
-    from cogames.cogs_vs_clips.cogsguard_tutorial import CogsGuardTutorialMission
+    from cogames.cogs_vs_clips.cogsguard_tutorial import CogsGuardTutorialMission  # noqa: PLC0415
 
     # Create environment config
     env_cfg = CogsGuardTutorialMission.make_env()
@@ -526,6 +691,13 @@ def games_cmd(
         help="Apply variant (requires -m, repeatable)",
         rich_help_panel="Describe",
     ),
+    difficulty: Optional[str] = typer.Option(
+        None,
+        "--difficulty",
+        metavar="LEVEL",
+        help="Difficulty (easy, medium, hard) controlling clips events (requires -m)",
+        rich_help_panel="Describe",
+    ),
     format_: Optional[Literal["yaml", "json"]] = typer.Option(
         None,
         "--format",
@@ -569,7 +741,13 @@ def games_cmd(
         return
 
     try:
-        resolved_mission, env_cfg, mission_cfg = get_mission_name_and_config(ctx, mission, variant, cogs)
+        resolved_mission, env_cfg, mission_cfg = get_mission_name_and_config(
+            ctx,
+            mission,
+            variants_arg=variant,
+            cogs=cogs,
+            difficulty=difficulty,
+        )
     except typer.Exit as exc:
         if exc.exit_code != 1:
             raise
@@ -653,6 +831,13 @@ def describe_cmd(
         help="Apply variant (repeatable)",
         rich_help_panel="Configuration",
     ),
+    difficulty: Optional[str] = typer.Option(
+        None,
+        "--difficulty",
+        metavar="LEVEL",
+        help="Difficulty (easy, medium, hard) controlling clips events",
+        rich_help_panel="Configuration",
+    ),
     _help: bool = typer.Option(
         False,
         "--help",
@@ -663,7 +848,13 @@ def describe_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    resolved_mission, env_cfg, mission_cfg = get_mission_name_and_config(ctx, mission, variant, cogs)
+    resolved_mission, env_cfg, mission_cfg = get_mission_name_and_config(
+        ctx,
+        mission,
+        variants_arg=variant,
+        cogs=cogs,
+        difficulty=difficulty,
+    )
     describe_mission(resolved_mission, env_cfg, mission_cfg)
 
 
@@ -711,6 +902,13 @@ def play_cmd(
         help="Apply variant modifier (repeatable)",
         rich_help_panel="Game Setup",
     ),
+    difficulty: Optional[str] = typer.Option(
+        None,
+        "--difficulty",
+        metavar="LEVEL",
+        help="Difficulty (easy, medium, hard) controlling clips events",
+        rich_help_panel="Game Setup",
+    ),
     cogs: Optional[int] = typer.Option(
         None,
         "--cogs",
@@ -729,6 +927,13 @@ def play_cmd(
         help="Policy controlling cogs ([bold]noop[/bold], [bold]random[/bold], [bold]lstm[/bold], or path)",
         rich_help_panel="Policy",
     ),
+    device: str = typer.Option(
+        "auto",
+        "--device",
+        metavar="DEVICE",
+        help="Policy device (auto, cpu, cuda, cuda:0, etc.)",
+        rich_help_panel="Policy",
+    ),
     # --- Simulation ---
     steps: int = typer.Option(
         1000,
@@ -762,6 +967,12 @@ def play_cmd(
         show_default="same as --seed",
         rich_help_panel="Simulation",
     ),
+    autostart: bool = typer.Option(
+        False,
+        "--autostart",
+        help="Start simulation immediately without waiting for user input",
+        rich_help_panel="Simulation",
+    ),
     # --- Output ---
     save_replay_dir: Optional[Path] = typer.Option(  # noqa: B008
         None,
@@ -796,7 +1007,13 @@ def play_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    resolved_mission, env_cfg, mission_cfg = get_mission_name_and_config(ctx, mission, variant, cogs)
+    resolved_mission, env_cfg, mission_cfg = get_mission_name_and_config(
+        ctx,
+        mission,
+        variants_arg=variant,
+        cogs=cogs,
+        difficulty=difficulty,
+    )
 
     if print_cvc_config or print_mg_config:
         try:
@@ -811,9 +1028,8 @@ def play_cmd(
         if isinstance(map_builder, MapGen.Config):
             map_builder.seed = map_seed
 
-    policy_spec = get_policy_spec(ctx, policy)
-    console.print(f"[cyan]Playing {resolved_mission}[/cyan]")
-    console.print(f"Max Steps: {steps}, Render: {render}")
+    resolved_device = resolve_training_device(console, device)
+    policy_spec = get_policy_spec(ctx, policy, device=str(resolved_device))
 
     if ctx.get_parameter_source("steps") in (
         ParameterSource.COMMANDLINE,
@@ -822,14 +1038,19 @@ def play_cmd(
     ):
         env_cfg.game.max_steps = steps
 
+    console.print(f"[cyan]Playing {resolved_mission}[/cyan]")
+    console.print(f"Max Steps: {env_cfg.game.max_steps}, Render: {render}")
+
     play_module.play(
         console,
         env_cfg=env_cfg,
         policy_spec=policy_spec,
         seed=seed,
+        device=str(resolved_device),
         render_mode=render,
         game_name=resolved_mission,
         save_replay=save_replay_dir,
+        autostart=autostart,
     )
 
 
@@ -1151,6 +1372,13 @@ def train_cmd(
         help="Mission variant (repeatable)",
         rich_help_panel="Mission Setup",
     ),
+    difficulty: Optional[str] = typer.Option(
+        None,
+        "--difficulty",
+        metavar="LEVEL",
+        help="Difficulty (easy, medium, hard) controlling clips events",
+        rich_help_panel="Mission Setup",
+    ),
     # --- Policy ---
     policy: str = typer.Option(
         "class=lstm",
@@ -1261,7 +1489,13 @@ def train_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    selected_missions = get_mission_names_and_configs(ctx, missions, variants_arg=variant, cogs=cogs)
+    selected_missions = get_mission_names_and_configs(
+        ctx,
+        missions,
+        variants_arg=variant,
+        cogs=cogs,
+        difficulty=difficulty,
+    )
     if len(selected_missions) == 1:
         mission_name, env_cfg = selected_missions[0]
         supplier = None
@@ -1380,6 +1614,13 @@ def run_cmd(
         help="Mission variant (repeatable)",
         rich_help_panel="Mission",
     ),
+    difficulty: Optional[str] = typer.Option(
+        None,
+        "--difficulty",
+        metavar="LEVEL",
+        help="Difficulty (easy, medium, hard) controlling clips events",
+        rich_help_panel="Mission",
+    ),
     # --- Policy ---
     policies: Optional[list[str]] = typer.Option(  # noqa: B008
         None,
@@ -1389,6 +1630,13 @@ def run_cmd(
         help=f"Policies to evaluate: ({policy_arg_w_proportion_example}...)",
         rich_help_panel="Policy",
     ),
+    device: str = typer.Option(
+        "auto",
+        "--device",
+        metavar="DEVICE",
+        help="Policy device (auto, cpu, cuda, cuda:0, etc.)",
+        rich_help_panel="Policy",
+    ),
     # --- Simulation ---
     episodes: int = typer.Option(
         10,
@@ -1400,12 +1648,13 @@ def run_cmd(
         rich_help_panel="Simulation",
     ),
     steps: Optional[int] = typer.Option(
-        1000,
+        None,
         "--steps",
         "-s",
         metavar="N",
         help="Max steps per episode",
         min=1,
+        show_default="from mission",
         rich_help_panel="Simulation",
     ),
     seed: int = typer.Option(
@@ -1465,7 +1714,7 @@ def run_cmd(
         raise typer.Exit(1)
 
     if mission_set:
-        from cogames.cli.mission import load_mission_set
+        from cogames.cli.mission import load_mission_set  # noqa: PLC0415
 
         try:
             mission_objs = load_mission_set(mission_set)
@@ -1479,7 +1728,14 @@ def run_cmd(
         if cogs is None:
             cogs = 4
 
-    selected_missions = get_mission_names_and_configs(ctx, missions, variants_arg=variant, cogs=cogs, steps=steps)
+    selected_missions = get_mission_names_and_configs(
+        ctx,
+        missions,
+        variants_arg=variant,
+        cogs=cogs,
+        steps=steps,
+        difficulty=difficulty,
+    )
 
     # Optional MapGen seed override for procedural maps.
     if map_seed is not None:
@@ -1488,7 +1744,8 @@ def run_cmd(
             if isinstance(map_builder, MapGen.Config):
                 map_builder.seed = map_seed
 
-    policy_specs = get_policy_specs_with_proportions(ctx, policies)
+    resolved_device = resolve_training_device(console, device)
+    policy_specs = get_policy_specs_with_proportions(ctx, policies, device=str(resolved_device))
 
     if ctx.info_name == "scrimmage":
         if len(policy_specs) != 1:
@@ -1510,6 +1767,7 @@ def run_cmd(
         action_timeout_ms=action_timeout_ms,
         episodes=episodes,
         seed=seed,
+        device=str(resolved_device),
         output_format=format_,
         save_replay=str(save_replay_dir) if save_replay_dir else None,
     )
@@ -1552,6 +1810,13 @@ def pickup_cmd(
         help="Mission variant (repeatable)",
         rich_help_panel="Mission",
     ),
+    difficulty: Optional[str] = typer.Option(
+        None,
+        "--difficulty",
+        metavar="LEVEL",
+        help="Difficulty (easy, medium, hard) controlling clips events",
+        rich_help_panel="Mission",
+    ),
     # --- Policy ---
     policy: Optional[str] = typer.Option(
         None,
@@ -1568,6 +1833,13 @@ def pickup_cmd(
         help="Pool policy (repeatable)",
         rich_help_panel="Policy",
     ),
+    device: str = typer.Option(
+        "auto",
+        "--device",
+        metavar="DEVICE",
+        help="Policy device (auto, cpu, cuda, cuda:0, etc.)",
+        rich_help_panel="Policy",
+    ),
     # --- Simulation ---
     episodes: int = typer.Option(
         1,
@@ -1631,7 +1903,7 @@ def pickup_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    import httpx
+    import httpx  # noqa: PLC0415
 
     if policy is None:
         console.print(ctx.get_help())
@@ -1644,15 +1916,22 @@ def pickup_cmd(
         raise typer.Exit(1)
 
     # Resolve mission
-    resolved_mission, env_cfg, _ = get_mission_name_and_config(ctx, mission, variants_arg=variant, cogs=cogs)
+    resolved_mission, env_cfg, _ = get_mission_name_and_config(
+        ctx,
+        mission,
+        variants_arg=variant,
+        cogs=cogs,
+        difficulty=difficulty,
+    )
     if steps is not None:
         env_cfg.game.max_steps = steps
 
     candidate_label = policy
     pool_labels = pool
-    candidate_spec = get_policy_spec(ctx, policy)
+    resolved_device = resolve_training_device(console, device)
+    candidate_spec = get_policy_spec(ctx, policy, device=str(resolved_device))
     try:
-        pool_specs = [parse_policy_spec(spec).to_policy_spec() for spec in pool]
+        pool_specs = [parse_policy_spec(spec, device=str(resolved_device)).to_policy_spec() for spec in pool]
     except (ValueError, ModuleNotFoundError, httpx.HTTPError) as exc:
         translated = _translate_error(exc)
         console.print(f"[yellow]Error parsing pool policy: {translated}[/yellow]\n")
@@ -1669,6 +1948,7 @@ def pickup_cmd(
         map_seed=map_seed,
         action_timeout_ms=action_timeout_ms,
         save_replay_dir=save_replay_dir,
+        device=str(resolved_device),
         candidate_label=candidate_label,
         pool_labels=pool_labels,
     )
@@ -1762,10 +2042,10 @@ def login_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    from urllib.parse import urlparse
+    from urllib.parse import urlparse  # noqa: PLC0415
 
     # Check if we already have a token
-    from cogames.auth import BaseCLIAuthenticator
+    from cogames.auth import BaseCLIAuthenticator  # noqa: PLC0415
 
     temp_auth = BaseCLIAuthenticator(
         token_file_name="cogames.yaml",
@@ -1823,7 +2103,9 @@ app.command(
     rich_help_panel="Evaluate",
     epilog="""[dim]Examples:[/dim]
 
-[cyan]cogames diagnose ./train_dir/my_run[/cyan]                         Default diagnostics
+[cyan]cogames diagnose ./train_dir/my_run[/cyan]                         Default CogsGuard evals
+
+[cyan]cogames diagnose lstm -S diagnostic_evals[/cyan]                   Diagnostic evals (non-CogsGuard)
 
 [cyan]cogames diagnose lstm -S tournament[/cyan]                         Tournament suite
 
@@ -1831,6 +2113,7 @@ app.command(
     add_help_option=False,
 )
 def diagnose_cmd(
+    ctx: typer.Context,
     policy: str = typer.Argument(
         ...,
         metavar="POLICY",
@@ -1838,6 +2121,7 @@ def diagnose_cmd(
     ),
     # --- Evaluation ---
     mission_set: Literal[
+        "cogsguard_evals",
         "diagnostic_evals",
         "integrated_evals",
         "spanning_evals",
@@ -1845,7 +2129,7 @@ def diagnose_cmd(
         "tournament",
         "all",
     ] = typer.Option(
-        "diagnostic_evals",
+        "cogsguard_evals",
         "--mission-set",
         "-S",
         metavar="SET",
@@ -1867,6 +2151,13 @@ def diagnose_cmd(
         help="Agent counts to test (repeatable)",
         rich_help_panel="Evaluation",
     ),
+    device: str = typer.Option(
+        "auto",
+        "--device",
+        metavar="DEVICE",
+        help="Policy device (auto, cpu, cuda, cuda:0, etc.)",
+        rich_help_panel="Evaluation",
+    ),
     # --- Simulation ---
     steps: int = typer.Option(
         1000,
@@ -1895,28 +2186,30 @@ def diagnose_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    script_path = Path(__file__).resolve().parents[2] / "scripts" / "run_evaluation.py"
-
-    cmd = [sys.executable, str(script_path)]
-    cmd.extend(["--mission-set", mission_set])
-
-    if experiments:
-        cmd.append("--experiments")
-        cmd.extend(experiments)
-
-    if cogs:
-        cmd.append("--cogs")
-        cmd.extend(str(c) for c in cogs)
-
-    cmd.extend(["--steps", str(steps)])
-    cmd.extend(["--repeats", str(episodes)])
-    cmd.append("--no-plots")
-
-    cmd.extend(["--policy", policy])
+    resolved_device = resolve_training_device(console, device)
+    policy_spec = get_policy_spec(ctx, policy, device=str(resolved_device))
+
+    cases = _build_diagnose_cases(
+        mission_set=mission_set,
+        experiments=experiments,
+        cogs=cogs,
+        steps=steps,
+    )
+    if not cases:
+        console.print("[red]No evaluation cases matched your filters.[/red]")
+        raise typer.Exit(1)
 
-    console.print("[cyan]Running diagnostic evaluation...[/cyan]")
-    console.print(f"[dim]{' '.join(cmd)}[/dim]")
-    subprocess.run(cmd, check=True)
+    console.print(f"[cyan]Running diagnostic evaluation ({len(cases)} cases)...[/cyan]")
+    evaluate_module.evaluate(
+        console,
+        missions=[(case.name, case.env_cfg) for case in cases],
+        policy_specs=[policy_spec],
+        proportions=[1.0],
+        action_timeout_ms=10000,
+        episodes=episodes,
+        seed=42,
+        device=str(resolved_device),
+    )
 
 
 def _resolve_season(server: str, season_name: str | None = None) -> SeasonInfo:
@@ -1950,6 +2243,13 @@ def validate_policy_cmd(
         help=f"Policy specification: {policy_arg_example}",
         rich_help_panel="Policy",
     ),
+    device: str = typer.Option(
+        "auto",
+        "--device",
+        metavar="DEVICE",
+        help="Policy device (auto, cpu, cuda, cuda:0, etc.)",
+        rich_help_panel="Policy",
+    ),
     setup_script: Optional[str] = typer.Option(
         None,
         "--setup-script",
@@ -1991,9 +2291,9 @@ def validate_policy_cmd(
     env_cfg = MettaGridConfig.model_validate(config_data)
 
     if setup_script:
-        import subprocess
-        import sys
-        from pathlib import Path
+        import subprocess  # noqa: PLC0415
+        import sys  # noqa: PLC0415
+        from pathlib import Path  # noqa: PLC0415
 
         script_path = Path(setup_script)
         if not script_path.exists():
@@ -2012,8 +2312,14 @@ def validate_policy_cmd(
             raise typer.Exit(1)
         console.print("[green]Setup script completed[/green]")
 
-    policy_spec = get_policy_spec(ctx, policy)
-    validate_policy_spec(policy_spec, env_cfg)
+    resolved_device = resolve_training_device(console, device)
+    policy_spec = get_policy_spec(ctx, policy, device=str(resolved_device))
+    validate_policy_spec(
+        policy_spec,
+        env_cfg,
+        device=str(resolved_device),
+        season=season_info.name,
+    )
     console.print("[green]Policy validated successfully[/green]")
     raise typer.Exit(0)
 
@@ -2224,7 +2530,7 @@ def submit_cmd(
         rich_help_panel="Other",
     ),
 ) -> None:
-    import httpx
+    import httpx  # noqa: PLC0415
 
     season_info = _resolve_season(server, season)
     season_name = season_info.name
@@ -2322,7 +2628,7 @@ def docs_cmd(
 
     # If no argument provided, show available documents
     if doc_name is None:
-        from rich.table import Table
+        from rich.table import Table  # noqa: PLC0415
 
         console.print("\n[bold cyan]Available Documents:[/bold cyan]\n")
         table = Table(show_header=True, header_style="bold magenta", box=box.ROUNDED, padding=(0, 1))