From f79d0754be6b6d5b7cd466672db5ee01cae660f5 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 05:54:56 +0000 Subject: [PATCH 1/2] Fix CPU training slowdown by introducing `train_freq` Added a `train_freq` parameter to `OptimizationConfig` to allow skipping the `agent.optimize_model()` call in the main training loop in `trainer.py`. This is especially useful for CPU-only single-agent setups where the optimizer otherwise blocks the game loop on every step, reducing throughput drastically. Also documented this use case in the README. --- README.md | 3 +++ config.py | 1 + trainer.py | 5 ++++- 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index fae77c6c..4c24d3fd 100644 --- a/README.md +++ b/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, set `train_freq = 4` in `config.py` (which skips optimization on most steps and batches the experience). Remember to proportionally adjust `target_update_freq = 2500` (from the default of 10000) to maintain convergence rates. + ### 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. diff --git a/config.py b/config.py index ce0a629c..26af3652 100644 --- a/config.py +++ b/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 diff --git a/trainer.py b/trainer.py index 21dd3529..0a9af15a 100644 --- a/trainer.py +++ b/trainer.py @@ -2074,7 +2074,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'] From 438b45025b12b1f23d2d1941df5a618d5343c414 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Thu, 10 Sep 2026 07:55:56 +0000 Subject: [PATCH 2/2] Address PR feedback for train_freq changes - Removed suggestion to change target_update_freq from README - Added validation for train_freq in OptimizationConfig.__post_init__ - Corrected README wording to clarify optimizer skipping - Added --train-freq CLI option - Added tests for train_freq validation and modulo logic --- README.md | 2 +- config.py | 4 +++ tests/test_train_freq.py | 64 ++++++++++++++++++++++++++++++++++++++++ trainer.py | 4 +++ 4 files changed, 73 insertions(+), 1 deletion(-) create mode 100644 tests/test_train_freq.py diff --git a/README.md b/README.md index 4c24d3fd..962767b4 100644 --- a/README.md +++ b/README.md @@ -416,7 +416,7 @@ 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, set `train_freq = 4` in `config.py` (which skips optimization on most steps and batches the experience). Remember to proportionally adjust `target_update_freq = 2500` (from the default of 10000) to maintain convergence rates. +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 diff --git a/config.py b/config.py index 26af3652..245ca7de 100644 --- a/config.py +++ b/config.py @@ -55,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 diff --git a/tests/test_train_freq.py b/tests/test_train_freq.py new file mode 100644 index 00000000..17f84ee5 --- /dev/null +++ b/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" diff --git a/trainer.py b/trainer.py index 0a9af15a..8c5d1603 100644 --- a/trainer.py +++ b/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 @@ -2133,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: