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

cogames/train.py

@@ -24,7 +24,7 @@ from mettagrid.policy.loader import (
 )
 from mettagrid.policy.policy import PolicySpec
 from mettagrid.policy.policy_env_interface import PolicyEnvInterface
-from mettagrid.policy.submission import POLICY_SPEC_FILENAME
+from mettagrid.policy.submission import POLICY_SPEC_FILENAME, SubmissionPolicySpec, write_submission_policy_spec
 from mettagrid.simulator import Simulator
 from mettagrid.util.stats_writer import NoopStatsWriter
 from pufferlib import pufferl
@@ -76,6 +76,25 @@ def _resolve_vector_counts(
     return num_envs, num_workers
 
 
+def _align_minibatch_size(minibatch_size: int, batch_size: int, bptt_horizon: int) -> int:
+    if bptt_horizon <= 1:
+        return minibatch_size
+
+    if minibatch_size < bptt_horizon:
+        return bptt_horizon
+
+    remainder = minibatch_size % bptt_horizon
+    if remainder == 0:
+        return minibatch_size
+
+    rounded_up = minibatch_size + (bptt_horizon - remainder)
+    if rounded_up <= batch_size:
+        return rounded_up
+
+    rounded_down = minibatch_size - remainder
+    return max(bptt_horizon, rounded_down)
+
+
 def train(
     env_cfg: Optional[MettaGridConfig],
     policy_class_path: str,
@@ -86,12 +105,14 @@ def train(
     seed: int,
     batch_size: int,
     minibatch_size: int,
+    map_seed: Optional[int] = None,
     missions_arg: Optional[list[str]] = None,
     vector_num_envs: Optional[int] = None,
     vector_batch_size: Optional[int] = None,
     vector_num_workers: Optional[int] = None,
     env_cfg_supplier: Optional[Callable[[], MettaGridConfig]] = None,
     log_outputs: bool = False,
+    checkpoint_interval: int = 200,
 ) -> None:
     console = Console()
 
@@ -101,7 +122,6 @@ def train(
     backend = pvector.Multiprocessing
     if platform.system() == "Darwin":
         multiprocessing.set_start_method("spawn", force=True)
-        backend = pvector.Serial
 
     try:
         cpu_cores = psutil.cpu_count(logical=False) or psutil.cpu_count(logical=True)
@@ -121,7 +141,7 @@ def train(
     else:
         num_workers = desired_workers
 
-    if backend is pvector.Multiprocessing and device.type != "cuda":
+    if backend is pvector.Multiprocessing and device.type not in ("cuda", "mps"):
         backend = pvector.Serial
         num_workers = 1
 
@@ -175,7 +195,7 @@ def train(
         envs_per_worker,
     )
 
-    env_creator = _EnvCreator(env_cfg, env_cfg_supplier)
+    env_creator = _EnvCreator(env_cfg, env_cfg_supplier, map_seed, seed)
     base_cfg = env_creator.clone_cfg()
 
     vecenv = pvector.make(
@@ -242,6 +262,16 @@ def train(
             amended_minibatch_size,
         )
 
+    aligned_minibatch_size = _align_minibatch_size(amended_minibatch_size, amended_batch_size, bptt_horizon)
+    if aligned_minibatch_size != amended_minibatch_size:
+        logger.info(
+            "Adjusting minibatch_size from %s to %s to align with bptt_horizon=%s",
+            amended_minibatch_size,
+            aligned_minibatch_size,
+            bptt_horizon,
+        )
+        amended_minibatch_size = aligned_minibatch_size
+
     effective_timesteps = max(num_steps, amended_batch_size)
     if effective_timesteps != num_steps:
         logger.info(
@@ -250,7 +280,6 @@ def train(
             effective_timesteps,
         )
 
-    checkpoint_interval = 200
     train_args = dict(
         env=env_name,
         device=device.type,
@@ -285,6 +314,7 @@ def train(
         vtrace_c_clip=1.0,
         prio_alpha=0.8,
         prio_beta0=0.2,
+        min_lr_ratio=0.0,
     )
 
     trainer = pufferl.PuffeRL(train_args, vecenv, policy.network())
@@ -354,13 +384,17 @@ def train(
             )
             console.print("=" * 80, style="bold green")
 
-        checkpoints = sorted(
-            {path.parent for path in checkpoints_path.rglob(POLICY_SPEC_FILENAME)},
-            key=lambda path: path.stat().st_mtime,
-        )
+        run_dir = checkpoints_path / trainer.logger.run_id
+        checkpoints = sorted(run_dir.glob("model_*.pt"), key=lambda path: path.stat().st_mtime)
 
         if checkpoints and not training_diverged:
             final_checkpoint = checkpoints[-1]
+            spec = SubmissionPolicySpec(
+                class_path=policy_class_path,
+                data_path=final_checkpoint.name,
+            )
+            write_submission_policy_spec(run_dir / POLICY_SPEC_FILENAME, spec)
+
             console.print()
             console.print(f"Final checkpoint: [cyan]{final_checkpoint}[/cyan]")
             if trainer.epoch < checkpoint_interval:
@@ -409,9 +443,13 @@ class _EnvCreator:
         self,
         env_cfg: Optional[MettaGridConfig],
         env_cfg_supplier: Optional[Callable[[], MettaGridConfig]],
+        map_seed: Optional[int],
+        fallback_seed: Optional[int],
     ) -> None:
         self._env_cfg = env_cfg
         self._env_cfg_supplier = env_cfg_supplier
+        self._map_seed = map_seed
+        self._fallback_seed = fallback_seed
 
     def clone_cfg(self) -> MettaGridConfig:
         if self._env_cfg_supplier is not None:
@@ -430,11 +468,11 @@ class _EnvCreator:
     ) -> PufferMettaGridEnv:
         target_cfg = cfg.model_copy(deep=True) if cfg is not None else self.clone_cfg()
 
-        # If this mission uses MapGen and the builder seed is unset, derive a deterministic
-        # MapGen seed from the per-env seed provided by the vectorized runner.
         map_builder = getattr(target_cfg.game, "map_builder", None)
-        if isinstance(map_builder, MapGen.Config) and seed is not None and map_builder.seed is None:
-            map_builder.seed = seed
+        if isinstance(map_builder, MapGen.Config):
+            base_seed = self._map_seed if self._map_seed is not None else self._fallback_seed
+            if base_seed is not None and (self._map_seed is not None or map_builder.seed is None):
+                map_builder.seed = base_seed + (seed or 0)
         simulator = Simulator()
         simulator.add_event_handler(StatsTracker(NoopStatsWriter()))
         simulator.add_event_handler(EarlyResetHandler())

cogames/verbose.py

@@ -4,14 +4,14 @@ import json
 
 from rich.console import Console
 
-from cogames.cogs_vs_clips.mission import Mission
+from cogames.cogs_vs_clips.mission import AnyMission
 from mettagrid.config.mettagrid_config import MettaGridConfig
 
 
 def print_configs(
     console: Console,
     env_cfg: MettaGridConfig,
-    mission_cfg: Mission | None = None,
+    mission_cfg: AnyMission | None = None,
     print_cvc_config: bool = False,
     print_mg_config: bool = False,
 ) -> None: