cogames 0.3.65__py3-none-any.whl → 0.3.69__py3-none-any.whl

cogames/cogs_vs_clips/reward_variants.py

@@ -0,0 +1,173 @@
+"""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
+
+import json
+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),
+        "aligner_lost": reward(stat("aligner.lost"), weight=-w_align_gear, max=-cap_align_gear),
+        "scrambler_gained": reward(stat("scrambler.gained"), weight=w_scramble_gear, max=cap_scramble_gear),
+        "scrambler_lost": reward(stat("scrambler.lost"), 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
+
+    # Parse JSON-encoded list strings (e.g., '["milestones"]' from sweeps)
+    if isinstance(variants, str):
+        if variants.startswith("["):
+            try:
+                parsed = json.loads(variants)
+                variant_names = list(parsed) if isinstance(parsed, list) else [variants]
+            except json.JSONDecodeError:
+                variant_names = [variants]
+        else:
+            variant_names = [variants]
+    else:
+        variant_names = 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/sites.py

@@ -72,16 +72,17 @@ MACHINA_1 = CoGameSite(
 
 def _cogsguard_hub_config() -> BaseHubConfig:
     return BaseHubConfig(
+        hub_object="c:hub",
         corner_bundle="extractors",
         cross_bundle="none",
         cross_distance=7,
         junction_object="junction",
+        heart_chest_object="c:chest",
         stations=[
-            "aligner_station",
-            "scrambler_station",
-            "miner_station",
-            "scout_station",
-            "chest",
+            "c:aligner",
+            "c:scrambler",
+            "c:miner",
+            "c:scout",
         ],
     )
 

cogames/cogs_vs_clips/stations.py

@@ -82,6 +82,7 @@ class CvCJunctionConfig(CvCStationConfig):
     def station_cfg(self, team: Optional[str] = None) -> GridObjectConfig:
         return GridObjectConfig(
             name="junction",
+            render_name="junction",
             render_symbol="📦",
             collective=team,
             aoes={
@@ -121,11 +122,12 @@ class CvCHubConfig(CvCStationConfig):
     attack_deltas: dict[str, int] = Field(default_factory=lambda: {"hp": -1, "influence": -100})
     elements: list[str] = Field(default_factory=lambda: CvCConfig.ELEMENTS)
 
-    def station_cfg(self, team: str) -> GridObjectConfig:
+    def station_cfg(self, team: str, collective: str | None = None) -> GridObjectConfig:
         return GridObjectConfig(
-            name="hub",
+            name=f"{team}:hub",
+            render_name="hub",
             render_symbol="📦",
-            collective=team,
+            collective=collective or team,
             aoes={
                 "influence": AOEConfig(
                     radius=self.aoe_range,
@@ -152,11 +154,12 @@ class CvCChestConfig(CvCStationConfig):
 
     heart_cost: dict[str, int] = Field(default_factory=lambda: CvCConfig.HEART_COST)
 
-    def station_cfg(self, team: str) -> GridObjectConfig:
+    def station_cfg(self, team: str, collective: str | None = None) -> GridObjectConfig:
         return GridObjectConfig(
-            name="chest",
+            name=f"{team}:chest",
+            render_name="chest",
             render_symbol="📦",
-            collective=team,
+            collective=collective or team,
             on_use_handlers={
                 "get_heart": Handler(
                     filters=[isAlignedToActor(), targetCollectiveHas({"heart": 1})],
@@ -180,12 +183,13 @@ class CvCGearStationConfig(CvCStationConfig):
     gear_costs: dict[str, dict[str, int]] = Field(default_factory=lambda: CvCConfig.GEAR_COSTS)
     gear_symbols: dict[str, str] = Field(default_factory=lambda: CvCConfig.GEAR_SYMBOLS)
 
-    def station_cfg(self, team: str) -> GridObjectConfig:
+    def station_cfg(self, team: str, collective: str | None = None) -> GridObjectConfig:
         cost = self.gear_costs.get(self.gear_type, {})
         return GridObjectConfig(
-            name=f"{self.gear_type}_station",
+            name=f"{team}:{self.gear_type}",
+            render_name=f"{self.gear_type}_station",
             render_symbol=self.gear_symbols[self.gear_type],
-            collective=team,
+            collective=collective or team,
             on_use_handlers={
                 "keep_gear": Handler(
                     filters=[isAlignedToActor(), actorHas({self.gear_type: 1})],

cogames/cogs_vs_clips/team.py

@@ -17,7 +17,8 @@ from mettagrid.config.mettagrid_config import (
 class CogTeam(Config):
     """Configuration for a cogs team."""
 
-    name: str = Field(default="cogs", description="Team name")
+    name: str = Field(default="cogs", description="Team name used for collectives and alignment")
+    short_name: str = Field(default="c", description="Short prefix used for map object names")
     wealth: int = Field(default=1, description="Wealth multiplier for initial resources")
     num_agents: int = Field(default=8, ge=1, description="Number of agents in the team")
 
@@ -28,6 +29,7 @@ class CogTeam(Config):
             CollectiveConfig with resource limits and initial inventory.
         """
         return CollectiveConfig(
+            name=self.name,
             inventory=InventoryConfig(
                 limits={
                     "resources": ResourceLimitsConfig(min=10000, resources=CvCConfig.ELEMENTS),

cogames/cogs_vs_clips/terrain.py

@@ -118,7 +118,7 @@ class MachinaArena(Scene[MachinaArenaConfig]):
 
         # Building weights
         default_building_weights = {
-            "chest": 0.0,
+            "c:chest": 0.0,
             "junction": 0.7,
             "germanium_extractor": 0.3,
             "silicon_extractor": 0.3,
@@ -377,7 +377,7 @@ class SequentialMachinaArena(Scene[SequentialMachinaArenaConfig]):
             raise ValueError(f"Unknown base_biome '{cfg.base_biome}'. Valid: {sorted(biome_map.keys())}")
         base_cfg: SceneConfig = BaseCfgModel.model_validate(cfg.base_biome_config or {})
         default_building_weights = {
-            "chest": 0.0,
+            "c:chest": 0.0,
             "junction": 0.6,
             "germanium_extractor": 0.2,
             "silicon_extractor": 0.2,

cogames/cogs_vs_clips/variants.py

@@ -2,9 +2,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, override
 
+from pydantic import Field
+
+from cogames.cogs_vs_clips.config import CvCConfig
 from cogames.cogs_vs_clips.evals.difficulty_variants import DIFFICULTY_VARIANTS
 from cogames.cogs_vs_clips.terrain import BaseHubVariant, MachinaArenaVariant
 from cogames.core import CoGameMissionVariant
+from mettagrid.config.mettagrid_config import MettaGridConfig
 from mettagrid.map_builder.map_builder import MapBuilderConfig
 from mettagrid.mapgen.mapgen import MapGen
 from mettagrid.mapgen.scenes.base_hub import DEFAULT_EXTRACTORS as HUB_EXTRACTORS
@@ -14,6 +18,37 @@ if TYPE_CHECKING:
     from cogames.cogs_vs_clips.mission import CvCMission
 
 
+def _apply_clips_settings(
+    mission: CvCMission,
+    *,
+    initial_clips_start: int | None = None,
+    initial_clips_spots: int | None = None,
+    scramble_start: int | None = None,
+    scramble_interval: int | None = None,
+    scramble_radius: int | None = None,
+    align_start: int | None = None,
+    align_interval: int | None = None,
+    align_radius: int | None = None,
+) -> None:
+    clips = mission.clips
+    if initial_clips_start is not None:
+        clips.initial_clips_start = initial_clips_start
+    if initial_clips_spots is not None:
+        clips.initial_clips_spots = initial_clips_spots
+    if scramble_start is not None:
+        clips.scramble_start = scramble_start
+    if scramble_interval is not None:
+        clips.scramble_interval = scramble_interval
+    if scramble_radius is not None:
+        clips.scramble_radius = scramble_radius
+    if align_start is not None:
+        clips.align_start = align_start
+    if align_interval is not None:
+        clips.align_interval = align_interval
+    if align_radius is not None:
+        clips.align_radius = align_radius
+
+
 class NumCogsVariant(CoGameMissionVariant):
     name: str = "num_cogs"
     description: str = "Set the number of cogs for the mission."
@@ -30,13 +65,91 @@ class NumCogsVariant(CoGameMissionVariant):
         mission.num_cogs = self.num_cogs
 
 
+class ClipsEasyVariant(CoGameMissionVariant):
+    name: str = "clips_easy"
+    description: str = "Slow clips expansion with late pressure."
+
+    @override
+    def modify_mission(self, mission: CvCMission) -> None:
+        _apply_clips_settings(
+            mission,
+            initial_clips_start=50,
+            initial_clips_spots=1,
+            scramble_start=250,
+            scramble_interval=250,
+            scramble_radius=15,
+            align_start=300,
+            align_interval=250,
+            align_radius=15,
+        )
+
+
+class ClipsMediumVariant(CoGameMissionVariant):
+    name: str = "clips_medium"
+    description: str = "Baseline clips pressure (Machina1 default)."
+
+    @override
+    def modify_mission(self, mission: CvCMission) -> None:
+        _apply_clips_settings(
+            mission,
+            initial_clips_start=10,
+            initial_clips_spots=1,
+            scramble_start=50,
+            scramble_interval=100,
+            scramble_radius=25,
+            align_start=100,
+            align_interval=100,
+            align_radius=25,
+        )
+
+
+class ClipsHardVariant(CoGameMissionVariant):
+    name: str = "clips_hard"
+    description: str = "Fast clips pressure with wider influence."
+
+    @override
+    def modify_mission(self, mission: CvCMission) -> None:
+        _apply_clips_settings(
+            mission,
+            initial_clips_start=5,
+            initial_clips_spots=2,
+            scramble_start=25,
+            scramble_interval=50,
+            scramble_radius=35,
+            align_start=50,
+            align_interval=50,
+            align_radius=35,
+        )
+
+
+class ClipsWaveOnlyVariant(CoGameMissionVariant):
+    name: str = "clips_wave_only"
+    description: str = "Initial clips wave only, no further spread."
+
+    @override
+    def modify_mission(self, mission: CvCMission) -> None:
+        disable_start = mission.max_steps + 1
+        _apply_clips_settings(
+            mission,
+            initial_clips_start=10,
+            initial_clips_spots=3,
+            scramble_start=disable_start,
+            scramble_interval=disable_start,
+            align_start=disable_start,
+            align_interval=disable_start,
+            scramble_radius=25,
+            align_radius=25,
+        )
+
+
 class DarkSideVariant(CoGameMissionVariant):
     name: str = "dark_side"
     description: str = "You're on the dark side of the asteroid. You recharge slower."
 
     @override
     def modify_mission(self, mission: CvCMission) -> None:
-        mission.cog.energy_regen = 0
+        mission.weather.day_deltas = {"solar": 0}
+        mission.weather.night_deltas = {"solar": 0}
 
 
 class SuperChargedVariant(CoGameMissionVariant):
@@ -45,7 +158,8 @@ class SuperChargedVariant(CoGameMissionVariant):
 
     @override
     def modify_mission(self, mission: CvCMission) -> None:
-        mission.cog.energy_regen += 2
+        mission.weather.day_deltas = {k: v + 2 for k, v in mission.weather.day_deltas.items()}
+        mission.weather.night_deltas = {k: v + 2 for k, v in mission.weather.night_deltas.items()}
 
 
 class EnergizedVariant(CoGameMissionVariant):
@@ -55,7 +169,8 @@ class EnergizedVariant(CoGameMissionVariant):
     @override
     def modify_mission(self, mission: CvCMission) -> None:
         mission.cog.energy_limit = max(mission.cog.energy_limit, 255)
-        mission.cog.energy_regen = mission.cog.energy_limit
+        mission.weather.day_deltas = {"solar": 255}
+        mission.weather.night_deltas = {"solar": 255}
 
 
 class Small50Variant(CoGameMissionVariant):
@@ -213,14 +328,65 @@ class BalancedCornersVariant(MachinaArenaVariant):
         node.max_balance_shortcuts = self.max_balance_shortcuts
 
 
+class MultiTeamVariant(CoGameMissionVariant):
+    """Split the map into multiple team instances, each with their own hub and resources."""
+
+    name: str = "multi_team"
+    description: str = "Split map into separate team instances with independent hubs."
+    num_teams: int = Field(default=2, ge=2, le=2, description="Number of teams (max 2 supported)")
+
+    @override
+    def modify_mission(self, mission: CvCMission) -> None:
+        team = next(iter(mission.teams.values()))
+        # Each team gets the original agent count; clear num_cogs so total is derived from teams
+        original_agents = mission.num_agents
+        mission.teams = {
+            name: team.model_copy(update={"name": name, "short_name": name, "num_agents": original_agents})
+            for name in ["cogs_green", "cogs_blue"][: self.num_teams]
+        }
+        mission.num_cogs = None
+
+    def modify_env(self, mission: CvCMission, env: MettaGridConfig) -> None:
+        original_builder = env.game.map_builder
+        # Shrink inner instance borders so teams are close together
+        if isinstance(original_builder, MapGen.Config):
+            original_builder.border_width = 1
+        env.game.map_builder = MapGen.Config(
+            instance=original_builder,
+            instances=self.num_teams,
+            set_team_by_instance=True,
+            instance_names=[t.short_name for t in mission.teams.values()],
+            instance_object_remap={
+                "c:hub": "{instance_name}:hub",
+                "c:chest": "{instance_name}:chest",
+                **{f"c:{g}": f"{{instance_name}}:{g}" for g in CvCConfig.GEAR},
+            },
+            # Connect instances: no added borders, clear walls at boundary
+            border_width=0,  # No outer border (inner instances have their own)
+            instance_border_width=0,  # No border between instances
+            instance_border_clear_radius=3,  # Clear walls near instance boundary
+        )
+
+
+class NoClipsVariant(CoGameMissionVariant):
+    name: str = "no_clips"
+    description: str = "Disable clips behavior entirely."
+
+    @override
+    def modify_mission(self, mission: CvCMission) -> None:
+        mission.clips.disabled = True
+
+
 VARIANTS: list[CoGameMissionVariant] = [
     CavesVariant(),
     CityVariant(),
     DarkSideVariant(),
+    NoClipsVariant(),
     DesertVariant(),
     EmptyBaseVariant(),
     EnergizedVariant(),
     ForestVariant(),
+    MultiTeamVariant(),
     QuadrantBuildingsVariant(),
     SingleResourceUniformVariant(),
     Small50Variant(),
@@ -228,4 +394,9 @@ VARIANTS: list[CoGameMissionVariant] = [
     *DIFFICULTY_VARIANTS,
 ]
 
-HIDDEN_VARIANTS: list[CoGameMissionVariant] = []
+HIDDEN_VARIANTS: list[CoGameMissionVariant] = [
+    ClipsEasyVariant(),
+    ClipsMediumVariant(),
+    ClipsHardVariant(),
+    ClipsWaveOnlyVariant(),
+]

cogames/cogs_vs_clips/weather.py

@@ -0,0 +1,52 @@
+"""Weather system events for CogsGuard missions.
+
+Day/night cycle that applies resource deltas to entities at regular intervals.
+"""
+
+from pydantic import Field
+
+from mettagrid.base_config import Config
+from mettagrid.config.event_config import EventConfig, periodic
+from mettagrid.config.mutation import updateTarget
+from mettagrid.config.tag import typeTag
+
+
+class WeatherConfig(Config):
+    """Configuration for day/night weather cycle."""
+
+    day_length: int = Field(default=200)
+    day_deltas: dict[str, int] = Field(default_factory=lambda: {"solar": 3})
+    night_deltas: dict[str, int] = Field(default_factory=lambda: {"solar": 1})
+    target_tag: str = Field(default="agent")
+
+    def events(self, max_steps: int) -> dict[str, EventConfig]:
+        """Create weather events for a mission.
+
+        Returns:
+            Dictionary of event name to EventConfig.
+        """
+        events: dict[str, EventConfig] = {}
+        tag = typeTag(self.target_tag)
+        half = self.day_length // 2
+
+        def _merge(apply: dict[str, int], reverse: dict[str, int]) -> dict[str, int]:
+            keys = set(apply) | set(reverse)
+            return {k: apply.get(k, 0) - reverse.get(k, 0) for k in keys}
+
+        # Dawn: reverse night deltas, apply day deltas
+        events["weather_day"] = EventConfig(
+            name="weather_day",
+            target_tag=tag,
+            timesteps=periodic(start=0, period=self.day_length, end=max_steps),
+            mutations=[updateTarget(_merge(self.day_deltas, self.night_deltas))],
+        )
+
+        # Dusk: reverse day deltas, apply night deltas
+        events["weather_night"] = EventConfig(
+            name="weather_night",
+            target_tag=tag,
+            timesteps=periodic(start=half, period=self.day_length, end=max_steps),
+            mutations=[updateTarget(_merge(self.night_deltas, self.day_deltas))],
+        )
+
+        return events

cogames/docs/SCRIPTED_AGENT.md

@@ -261,11 +261,11 @@ uv run cogames play --mission evals.diagnostic_assemble_seeded_search -p baselin
 
 ```bash
 # Run full evaluation suite
-uv run python packages/cogames/scripts/run_evaluation.py
+uv run cogames diagnose ladybug -S all
 
 # Evaluate specific agent
-uv run python packages/cogames/scripts/run_evaluation.py --policy baseline
-uv run python packages/cogames/scripts/run_evaluation.py --policy ladybug
+uv run cogames diagnose baseline
+uv run cogames diagnose ladybug
 ```
 
 ## Evaluation Results

cogames/evaluate.py

@@ -13,10 +13,10 @@ from pydantic import BaseModel, ConfigDict
 from rich.console import Console
 from rich.table import Table
 
-from metta_alo.rollout import run_multi_episode_rollout
 from metta_alo.scoring import allocate_counts, validate_proportions
 from mettagrid import MettaGridConfig
 from mettagrid.policy.policy import PolicySpec
+from mettagrid.runner.rollout import run_multi_episode_rollout
 from mettagrid.simulator.multi_episode.rollout import MultiEpisodeRolloutResult
 from mettagrid.simulator.multi_episode.summary import MultiEpisodeRolloutSummary, build_multi_episode_rollout_summaries
 
@@ -46,6 +46,7 @@ def evaluate(
     episodes: int,
     action_timeout_ms: int,
     seed: int = 42,
+    device: Optional[str] = None,
     output_format: Optional[Literal["yaml", "json"]] = None,
     save_replay: Optional[str] = None,
 ) -> MissionResultsSummary:
@@ -70,7 +71,7 @@ def evaluate(
     all_replay_paths: list[str] = []
     for mission_name, env_cfg in missions:
         counts = allocate_counts(env_cfg.game.num_agents, proportions)
-        assignments = np.repeat(np.arange(len(counts), dtype=int), counts)
+        assignments = [i for i, c in enumerate(counts) for _ in range(c)]
 
         progress_label = f"Simulating ({mission_name})"
         with typer.progressbar(length=episodes, label=progress_label) as progress:
@@ -83,6 +84,7 @@ def evaluate(
                 max_action_time_ms=action_timeout_ms,
                 replay_dir=save_replay,
                 create_replay_dir=save_replay is not None,
+                device=device,
                 on_progress=lambda _episode_idx, _result: progress.update(1),
             )