Merge pull request #400 from relh/richard-cogames-train

Add cogames.cogs-v-clips install/training command

Author: jsuarez5341

pufferlib/config/cogames.ini

@@ -0,0 +1,21 @@
+[base]
+package = cogames
+env_name = cogames.cogs_v_clips.training_facility.harvest cogames.cogs_v_clips.training_facility.assemble cogames.cogs_v_clips.machina_1.open_world
+policy_name = Policy
+rnn_name = Recurrent
+
+[vec]
+num_envs = 64
+num_workers = 16
+batch_size = auto
+zero_copy = True
+
+[env]
+render_mode = none
+variants = heart_chorus inventory_heart_tune
+
+[train]
+total_timesteps = 50_000_000
+batch_size = auto
+minibatch_size = 1024
+bptt_horizon = 64 

pufferlib/environments/cogames/__init__.py

@@ -0,0 +1,9 @@
+"""CoGames integration package."""
+
+from .environment import env_creator, make
+
+try:
+    import torch
+    from .torch import Policy, Recurrent
+except ImportError:
+    pass

pufferlib/environments/cogames/environment.py

@@ -0,0 +1,27 @@
+"""CoGames wrapper for PufferLib."""
+
+import functools
+from cogames.cli.mission import get_mission
+from mettagrid import PufferMettaGridEnv
+from mettagrid.envs.stats_tracker import StatsTracker
+from mettagrid.simulator import Simulator
+from mettagrid.util.stats_writer import NoopStatsWriter
+
+
+def env_creator(name="cogames.cogs_v_clips.machina_1.open_world"):
+    return functools.partial(make, name=name)
+
+
+def make(name="cogames.cogs_v_clips.machina_1.open_world", variants=None, cogs=None, render_mode="auto", seed=None, buf=None):
+    mission_name = name.removeprefix("cogames.cogs_v_clips.") if name.startswith("cogames.cogs_v_clips.") else name
+    variants = variants.split() if isinstance(variants, str) else variants
+    _, env_cfg, _ = get_mission(mission_name, variants_arg=variants, cogs=cogs)
+
+    render = "none" if render_mode == "auto" else "unicode" if render_mode in {"human", "ansi"} else render_mode
+    simulator = Simulator()
+    simulator.add_event_handler(StatsTracker(NoopStatsWriter()))
+    env = PufferMettaGridEnv(simulator=simulator, cfg=env_cfg, buf=buf, seed=seed or 0)
+    env.render_mode = render
+    if seed:
+        env.reset(seed)
+    return env

pufferlib/environments/cogames/torch.py

@@ -0,0 +1,25 @@
+"""Torch policies for CoGames environments."""
+
+import torch
+import pufferlib.models
+import pufferlib.pytorch
+
+
+class Policy(pufferlib.models.Default):
+    def __init__(self, env, hidden_size: int = 256, **kwargs):
+        super().__init__(env, hidden_size=hidden_size)
+        self.register_buffer("_inv_scale", torch.tensor(1.0 / 255.0), persistent=False)
+
+    def encode_observations(self, observations, state=None):
+        batch_size = observations.shape[0]
+        if self.is_dict_obs:
+            obs_map = pufferlib.pytorch.nativize_tensor(observations, self.dtype)
+            flattened = torch.cat([v.view(batch_size, -1) for v in obs_map.values()], dim=1)
+        else:
+            flattened = observations.view(batch_size, -1).float() * self._inv_scale
+        return self.encoder(flattened)
+
+
+class Recurrent(pufferlib.models.LSTMWrapper):
+    def __init__(self, env, policy, input_size: int = 256, hidden_size: int = 256):
+        super().__init__(env, policy, input_size=input_size, hidden_size=hidden_size)

pyproject.toml

@@ -121,6 +121,17 @@ metta = [
     'metta-mettagrid @ git+https://github.com/metta-ai/metta.git@main#subdirectory=mettagrid',
 ]
 
+cogames = [
+    'gym',
+    'gymnasium',
+    'omegaconf',
+    'hydra-core',
+    'duckdb',
+    'raylib>=5.5.0',
+    'mettagrid @ git+https://github.com/metta-ai/metta.git@main#subdirectory=packages/mettagrid',
+    'cogames @ git+https://github.com/metta-ai/metta.git@main#subdirectory=packages/cogames',
+]
+
 microrts = [
     'gym==0.23',
     'gymnasium==0.29.1',

setup.py

@@ -275,8 +275,8 @@ def run(self):
     'numpy<2.0',
     'shimmy[gym-v21]',
     'gym==0.23',
-    'gymnasium==0.29.1',
-    'pettingzoo==1.24.1',
+    'gymnasium>=0.29.1',
+    'pettingzoo>=1.24.1',
 ]
 
 if not NO_TRAIN: