Fix CPU training slowdown by introducing train_freq

README.md

@@ -415,6 +415,9 @@ Multiple Chrome instances run in parallel via `SubprocVecEnv` (multiprocessing).
 
 Auto-scaling monitors CPU, RAM, and step latency to add or remove agents dynamically. Each agent uses ~500MB RAM.
 
+#### CPU Single-Agent Training
+When training on a CPU with a single agent, the default optimization configuration may block the game loop and severely reduce the training step rate. To improve throughput, you can run the trainer with `--train-freq 4`. Setting `train_freq = 4` runs one optimization step for every four environment steps, reducing CPU load while continuing to add every collected transition to the replay buffer.
+
 ### Experience Replay
 
 We use **Prioritized Experience Replay (PER)** with a SumTree data structure. Transitions with high TD-error (where the network's prediction was most wrong) get sampled more frequently. This means the network spends more time learning from surprising or difficult situations.

config.py

@@ -31,6 +31,7 @@ class OptimizationConfig:
     eps_start: float = 1.0
     eps_end: float = 0.08          # Less randomness at convergence
     eps_decay: int = 8000          # Calibrated for steps_done += 1 per batch
+    train_freq: int = 1            # Frequency of optimization steps (1 for GPU, 4+ recommended for CPU)
     target_update_freq: int = 10000 # Increased for stability with 10 agents
     max_episodes: int = 5000000
     checkpoint_every: int = 50
@@ -54,6 +55,10 @@ class OptimizationConfig:
     super_pattern_straight_penalty_cap: float = 0.1
     super_pattern_food_reward_cap: float = 15.0
 
+    def __post_init__(self):
+        if self.train_freq < 1:
+            raise ValueError(f"train_freq must be >= 1, got {self.train_freq}")
+
 @dataclass
 class ReplayBufferConfig:
     capacity: int = 100000

tests/test_train_freq.py

@@ -0,0 +1,64 @@
+import pytest
+import sys
+import os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from config import OptimizationConfig
+
+def test_train_freq_validation():
+    # Should be valid
+    cfg = OptimizationConfig(train_freq=1)
+    assert cfg.train_freq == 1
+
+    cfg = OptimizationConfig(train_freq=4)
+    assert cfg.train_freq == 4
+
+    # Should raise ValueError
+    with pytest.raises(ValueError):
+        OptimizationConfig(train_freq=0)
+
+    with pytest.raises(ValueError):
+        OptimizationConfig(train_freq=-1)
+
+def test_train_freq_modulo_logic():
+    # Simulate trainer logic
+    class DummyAgent:
+        def __init__(self):
+            self.calls = 0
+        def optimize_model(self):
+            self.calls += 1
+            return {"loss": 0.1}
+
+    agent = DummyAgent()
+
+    # Simulate train_freq=1
+    train_freq = 1
+    total_steps = 0
+    for _ in range(10):
+        total_steps += 1
+        if total_steps % train_freq == 0:
+            agent.optimize_model()
+
+    assert agent.calls == 10, "With train_freq=1, optimize_model should be called on every step"
+
+    # Simulate train_freq=4
+    agent.calls = 0
+    train_freq = 4
+    total_steps = 0
+    for _ in range(10):
+        total_steps += 1
+        if total_steps % train_freq == 0:
+            agent.optimize_model()
+
+    assert agent.calls == 2, "With train_freq=4, optimize_model should be called twice in 10 steps (at 4 and 8)"
+
+    # Verify target update freq is independent
+    target_updates = 0
+    target_update_freq = 10000
+    total_steps = 0
+    for _ in range(20000):
+        total_steps += 1
+        if total_steps % target_update_freq == 0:
+            target_updates += 1
+
+    assert target_updates == 2, "target updates should happen exactly twice in 20000 steps"

trainer.py

@@ -1446,6 +1446,9 @@ def train(args):
     if args.vision_size:
         cfg.env.resolution = (args.vision_size, args.vision_size)
 
+    if args.train_freq > 0:
+        cfg.opt.train_freq = args.train_freq
+
     # Backend selection
     cfg.browser_backend = args.backend
     cfg.ws_server_url = args.ws_server_url
@@ -2074,7 +2077,10 @@ def finalize_episode(agent_index, terminal_state, cause, force_done_flag):
                         dashboard.log_event(f"Scale DOWN -> {env.num_agents} agents")
 
             # Train
-            metrics = agent.optimize_model()
+            metrics = None
+            if total_steps % cfg.opt.train_freq == 0:
+                metrics = agent.optimize_model()
+
             if metrics is not None:
                 last_metrics = metrics
                 train._last_loss = metrics['loss']
@@ -2130,6 +2136,7 @@ def finalize_episode(agent_index, terminal_state, cause, force_done_flag):
     parser.add_argument("--ai-lookback", type=int, default=500, help="AI Supervisor: analyze last N episodes (default: 500)")
     parser.add_argument("--ai-model", type=str, default=None, help="AI Supervisor: override LLM model name")
     parser.add_argument("--ai-key", type=str, default=None, help="AI Supervisor: API key (default: from env var)")
+    parser.add_argument("--train-freq", type=int, default=0, help="Override OptimizationConfig train_freq (e.g. 4 for CPU)")
     args = parser.parse_args()
 
     if args.reset: