perf(evolution): cut per-step dead work + add CMA-ES diagonal mode opt-in - #133
Merged
Conversation
…t-in
Five changes that together unblock LSTMPPO+klinotaxis evolution campaigns
that were previously intractable on the order of hours-per-generation.
Per-step path fixes (small but real, compound across 1000-step episodes):
- evolution/fitness.py: pass `theme=Theme.HEADLESS` to
`create_env_from_config` so the per-step `_render_step` short-circuits
in worker processes (which have no terminal). Was defaulting to
`Theme.ASCII`, building the grid and `print()`ing every step.
- evolution/loop.py `_init_worker`: add `torch.set_num_threads(1)` and
silence the `quantumnematode.agent.runners` and
`quantumnematode.agent.agent` loggers to WARNING. Stops BLAS
oversubscription at parallel_workers > 1; skips per-step logger.info
f-string construction without changing behaviour.
- agent/runners.py: gate per-step `logger.info` and `logger.debug`
f-strings with `isEnabledFor(...)` so the f-string isn't built when
the level is filtered.
The big one — CMA-ES diagonal mode opt-in:
- optimizers/evolutionary.py: add `diagonal: bool = False` kwarg on
CMAESOptimizer.__init__. When True, sets `CMA_diagonal=True` in the
underlying cma library, restricting the covariance matrix to its
diagonal. Drops `tell()` cost from O(n²) to O(n) — a tractability
requirement for genome dim >~1000.
- utils/config_loader.py: add `cma_diagonal: bool = False` field on
EvolutionConfig (back-compat default), with a comment explaining the
scaling rationale.
- scripts/run_evolution.py: plumb `evolution_config.cma_diagonal`
through to CMAESOptimizer.
- configs/evolution/lstmppo_foraging_small_klinotaxis.yml: opt in
(`cma_diagonal: true`) so the existing LSTMPPO smoke pilot is
actually runnable.
Profile data (LSTMPPO weight genome, dim=46989):
Without diagonal: With diagonal:
init: 10.6 s init: 0.003 s
ask(): 3-8 s ask(): 0.002 s
tell(): 167 s tell(): 0.003 s
End-to-end bench (population 4, 2 episodes, 1 generation):
Before: >5 minutes (timed out, harness-killed before completing)
After: 0.80 s (0.100 s per episode)
The per-step fixes alone wouldn't move the needle on LSTMPPO — full-cov
CMA-ES `tell()` at 47k dims dominates everything else. The diagonal
opt-in is the change that makes M2 LSTMPPO arms feasible.
Bench harness committed at scripts/benchmarks/bench_evolution_smoke.py
with line-buffered stdout (so partial progress is visible if the run is
killed mid-bench) and a `--cma-diagonal` flag for ad-hoc comparison.
Tests added (3): test_init_worker_sets_perf_policy locks in the worker
contract; test_cmaes_diagonal_full_loop verifies the diagonal mode works
end-to-end; test_cmaes_diagonal_default_off proves back-compat.
Verification: pre-commit run -a clean (10 hooks); pytest -m "not nightly"
2191 passed in 81 s; manual bench numbers above.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Spell out the convergence trade-off where users will see it: - EvolutionConfig.cma_diagonal field comment: explain that diagonal mode needs typically 2-10x more generations to reach the same fitness on non-separable problems (Ros & Hansen 2008), but full-cov isn't a competing option at large n, so net wall-clock is dramatically faster with diagonal anyway. Note when to keep it False (small genomes). - CMAESOptimizer.diagonal docstring: same trade-off + sep-CMA-ES proper name + size guidance. - LSTMPPO pilot config comment: ack the per-generation slowdown so future readers know why cma_diagonal: true is set there. Update the nematode-run-evolution skill to reflect post-perf-fix reality: - Replaced the "LSTMPPO is slow, always background, takes 1-3 min for smoke" narrative with the actual cause (cma_diagonal=False on a weight-evolution genome). Smoke runs are now sub-second on either brain when the YAML has cma_diagonal: true. - Added a "When to enable cma_diagonal" section with a genome-dim decision table (n<100 false; 100-1000 prefer true; n>1000 mandatory). - Updated runtime examples: LSTMPPO 10gen × pop 8 × 3ep is ~24 s, not ~3 hours. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Address review findings before opening the PR.
Blocking fix — `_init_worker` was silencing logger names that don't
exist:
- `runners.py:22` and `agent.py:23` both do `from
quantumnematode.logging_config import logger`, so the actual logger
used at the per-step `isEnabledFor(...)` gates is named
`quantumnematode.logging_config` — NOT `quantumnematode.agent.runners`
or `quantumnematode.agent.agent`.
- Our `_init_worker` was setting WARNING on the latter two phantom
names, leaving the real logger at whatever level the parent had it
(INFO under `--log-level INFO`). At `parallel_workers > 1` with INFO
logging, the per-step f-string skip in `runners.py` would silently
not fire, defeating the gate.
- Fix: silence `quantumnematode.logging_config` instead. The loop's
own module logger (`quantumnematode.evolution.loop`) keeps the
parent's verbosity so generation-level progress still surfaces.
Test that would have caught the bug — `test_init_worker_sets_perf_policy`
now imports the actual runtime logger via `from
quantumnematode.agent.runners import logger as runtime_logger` and
asserts `not runtime_logger.isEnabledFor(logging.INFO)` after init.
The previous assertion (level on phantom names) passed regardless of
whether the real logger was affected — exactly the bug class we
should have caught. Test also forces root to INFO + runtime logger
to NOTSET first, so the global pytest WARNING level can't mask the
bug.
End-to-end YAML→optimiser regression — added
`test_cma_diagonal_yaml_propagates_to_optimizer_options` that loads a
`cma_diagonal: true` YAML and asserts the cma library's
`opts['CMA_diagonal']` is True after constructing `CMAESOptimizer`.
Also added `test_lstmppo_pilot_config_enables_cma_diagonal` to lock in
that the shipped LSTMPPO pilot keeps `cma_diagonal: true` (regressing
that would silently make the pilot unrunnable).
Test robustness — `test_cmaes_diagonal_default_off` now asserts `not
opts.get("CMA_diagonal")` instead of `== 0`. cma's "off" sentinel is
0/0.0 today but the library reserves the right to change it.
Stale docs — `scripts/run_evolution.py` Timing section now reflects
post-perf-fix reality (~50 ms/episode MLPPPO, ~100 ms/episode LSTMPPO)
and points at `cma_diagonal` as the real gotcha for large genomes.
Config explanation — `configs/evolution/mlpppo_foraging_small.yml`
gains a comment explaining why `cma_diagonal: false` is preserved
(byte-identical with the M0 framework smoke test, which other tooling
relies on). Notes that real M2-scale MLPPPO weight campaigns should
opt in.
Verification: pre-commit run -a clean (10 hooks); pytest -m "not
nightly" 2193 passed in 83 s (was 2191; added 2 new tests).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdds a configurable diagonal CMA-ES mode and wires it through configs, optimizer, CLI/bench scripts, and tests; limits PyTorch threads and suppresses worker INFO logging in multiprocessing workers; forces headless envs for fitness evaluation; and guards expensive per-step logging. Docs and timing guidance updated. Changes
Sequence Diagram(s)sequenceDiagram
participant CLI
participant ConfigLoader
participant EvolutionLoop
participant Optimizer
participant CMA_LIB
participant EnvFactory
CLI->>ConfigLoader: load YAML (includes cma_diagonal)
ConfigLoader->>EvolutionLoop: provide SimulationConfig + EvolutionConfig
EvolutionLoop->>Optimizer: init(genome_dim, diagonal=cma_diagonal)
Optimizer->>CMA_LIB: create CMAEvolutionStrategy(opts including CMA_diagonal)
EvolutionLoop->>EnvFactory: create_env_from_config(theme=HEADLESS, seed)
EvolutionLoop->>Optimizer: ask() -> population
EvolutionLoop->>EnvFactory: evaluate(population members)
EvolutionLoop->>Optimizer: tell(fitnesses)
Optimizer->>CMA_LIB: update (tell -> update covariance)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
chrisjz
marked this pull request as ready for review
April 27, 2026 10:31
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
packages/quantum-nematode/tests/quantumnematode_tests/optimizers/test_evolutionary.py (1)
153-188: Add explicit return type annotations to the new test methods.Please annotate these test methods with
-> Noneto match repository typing rules.Proposed patch
- def test_cmaes_diagonal_full_loop(self): + def test_cmaes_diagonal_full_loop(self) -> None: @@ - def test_cmaes_diagonal_default_off(self): + def test_cmaes_diagonal_default_off(self) -> None:As per coding guidelines "Use comprehensive type annotations in all code".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@packages/quantum-nematode/tests/quantumnematode_tests/optimizers/test_evolutionary.py` around lines 153 - 188, The two new test functions test_cmaes_diagonal_full_loop and test_cmaes_diagonal_default_off lack explicit return type annotations; update their definitions to include "-> None" (e.g., def test_cmaes_diagonal_full_loop(...) -> None:) to comply with repository typing rules and ensure any linters/type-checkers accept the change.scripts/benchmarks/bench_evolution_smoke.py (1)
140-144:.bench_evolution_tmpis already gitignored.The directory is properly configured in
.gitignore, so commit-tracked artifacts won't accumulate in the repository. Consider usingtempfile.TemporaryDirectory()if automatic cleanup is desired, but this is not required given the existing gitignore entry.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/benchmarks/bench_evolution_smoke.py` around lines 140 - 144, The comment claiming ".bench_evolution_tmp" prevents repo pollution is misleading because that directory is already listed in .gitignore; either remove/update the comment to reflect the gitignore entry or switch to using a transient temp directory via tempfile.TemporaryDirectory() to get automatic cleanup. Locate the tmp_dir variable and its creation (tmp_dir = PROJECT_ROOT / ".bench_evolution_tmp" and tmp_dir.mkdir(exist_ok=True)) in bench_evolution_smoke.py and either (a) remove the misleading comment and leave the current creation as-is, or (b) replace the PROJECT_ROOT-based directory with a tempfile.TemporaryDirectory() context manager and use its path for outputs to ensure automatic cleanup.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In
`@packages/quantum-nematode/tests/quantumnematode_tests/optimizers/test_evolutionary.py`:
- Around line 153-188: The two new test functions test_cmaes_diagonal_full_loop
and test_cmaes_diagonal_default_off lack explicit return type annotations;
update their definitions to include "-> None" (e.g., def
test_cmaes_diagonal_full_loop(...) -> None:) to comply with repository typing
rules and ensure any linters/type-checkers accept the change.
In `@scripts/benchmarks/bench_evolution_smoke.py`:
- Around line 140-144: The comment claiming ".bench_evolution_tmp" prevents repo
pollution is misleading because that directory is already listed in .gitignore;
either remove/update the comment to reflect the gitignore entry or switch to
using a transient temp directory via tempfile.TemporaryDirectory() to get
automatic cleanup. Locate the tmp_dir variable and its creation (tmp_dir =
PROJECT_ROOT / ".bench_evolution_tmp" and tmp_dir.mkdir(exist_ok=True)) in
bench_evolution_smoke.py and either (a) remove the misleading comment and leave
the current creation as-is, or (b) replace the PROJECT_ROOT-based directory with
a tempfile.TemporaryDirectory() context manager and use its path for outputs to
ensure automatic cleanup.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 06a95777-5cfe-4338-ae65-85f2c3414cf8
📒 Files selected for processing (16)
.claude/skills/nematode-run-evolution/skill.md.gitignoreconfigs/evolution/lstmppo_foraging_small_klinotaxis.ymlconfigs/evolution/mlpppo_foraging_small.ymlpackages/quantum-nematode/quantumnematode/agent/runners.pypackages/quantum-nematode/quantumnematode/evolution/fitness.pypackages/quantum-nematode/quantumnematode/evolution/loop.pypackages/quantum-nematode/quantumnematode/optimizers/evolutionary.pypackages/quantum-nematode/quantumnematode/utils/config_loader.pypackages/quantum-nematode/tests/quantumnematode_tests/evolution/test_config.pypackages/quantum-nematode/tests/quantumnematode_tests/evolution/test_fitness.pypackages/quantum-nematode/tests/quantumnematode_tests/evolution/test_loop_smoke.pypackages/quantum-nematode/tests/quantumnematode_tests/optimizers/test_evolutionary.pyscripts/benchmarks/__init__.pyscripts/benchmarks/bench_evolution_smoke.pyscripts/run_evolution.py
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
The previous comment ("so we don't pollute repo state") was a partial
explanation: it captured the gitignore aspect but missed that we keep
the dir project-local rather than under /tmp so artefacts survive for
inspection across invocations. Updated to spell both reasons out and
to explicitly contrast with tempfile.TemporaryDirectory() (which would
auto-delete and defeat the inspection use case).
Also reviewer raised: should we switch to tempfile.TemporaryDirectory()
for automatic cleanup? Rejected — the bench harness is for inspecting
loop output, so auto-deletion would be a workflow regression.
Other reviewer nit (add ``-> None`` to the two new test methods in
test_evolutionary.py) NOT applied: the file's pre-existing 9 test
methods all omit return annotations, ruff and pyright both pass clean
on the file as-is, so adding ``-> None`` to just 2 methods would
introduce gratuitous inconsistency for no checker-driven benefit.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
scripts/benchmarks/bench_evolution_smoke.py (1)
43-68: Align function docstrings with NumPy style (and add one for_parse_args).
_parse_args()has no docstring, andmain()uses a one-liner. Bringing both to NumPy-style keeps this file consistent with repo docstring rules.Proposed fix
def _parse_args() -> argparse.Namespace: + """Parse CLI arguments for the evolution smoke benchmark. + + Returns + ------- + argparse.Namespace + Parsed command-line arguments. + """ parser = argparse.ArgumentParser(description=__doc__) @@ def main() -> int: - """Run a one-generation benchmark and print wall-clock numbers.""" + """Run a one-generation benchmark. + + Returns + ------- + int + Process exit code (0 on success, 1 on validation/config errors). + """As per coding guidelines,
**/*.py: "Use NumPy-style docstrings".🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/benchmarks/bench_evolution_smoke.py` around lines 43 - 68, Add NumPy-style docstrings for both functions: write a short summary line, followed by blank line and sections as appropriate. For _parse_args(), document the purpose, parameters (none) and Returns (argparse.Namespace) describing the parsed CLI options. For main(), expand the one-liner into a NumPy-style docstring that describes what the function does, the return type (int), and any side effects (prints wall-clock numbers). Place the docstrings directly above the corresponding function definitions (_parse_args and main) using the NumPy conventions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@scripts/benchmarks/bench_evolution_smoke.py`:
- Around line 69-70: Add explicit lower-bound checks in main() after calling
_parse_args(): verify args.population, args.episodes, and args.parallel are all
greater than 0 and fail fast with a clear error (raise ValueError or call
sys.exit with a descriptive message) if any are <= 0; update the validation to
reference args.population, args.episodes, and args.parallel so invalid benchmark
sizing is rejected before proceeding to the optimizer/run logic.
---
Nitpick comments:
In `@scripts/benchmarks/bench_evolution_smoke.py`:
- Around line 43-68: Add NumPy-style docstrings for both functions: write a
short summary line, followed by blank line and sections as appropriate. For
_parse_args(), document the purpose, parameters (none) and Returns
(argparse.Namespace) describing the parsed CLI options. For main(), expand the
one-liner into a NumPy-style docstring that describes what the function does,
the return type (int), and any side effects (prints wall-clock numbers). Place
the docstrings directly above the corresponding function definitions
(_parse_args and main) using the NumPy conventions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a611446e-0f29-4be8-be54-ab67475d5896
📒 Files selected for processing (1)
scripts/benchmarks/bench_evolution_smoke.py
\`--population 0\`, \`--episodes 0\`, or \`--parallel 0\` previously
surfaced as a confusing cma library error mid-run ("number of weights
must be >=2, was 0"), after partial setup output had already printed.
Pydantic's Field(ge=1) constraints on EvolutionConfig DO protect the
loop, but the bench's \`base_evolution.model_copy(update={...})\` path
bypasses those constraints in Pydantic v2 (model_copy doesn't
revalidate by default).
Add explicit lower-bound checks at the CLI boundary so invalid sizing
prints a clear error and exits with code 1 before any optimiser code
runs.
Reviewer also asked for NumPy-style docstrings on \`_parse_args\` and
\`main\`. Not applied: the canonical CLI in scripts/run_evolution.py
uses single-line docstrings on its main() ("Entry point.") and
parse_arguments() ("Parse command-line arguments."); the bench's
existing one-line docstrings are already richer than that convention.
Adding multi-section NumPy docstrings to private CLI helpers would
violate the repo style and CLAUDE-instruction against multi-paragraph
docstrings.
Verified: \`uv run python scripts/benchmarks/bench_evolution_smoke.py
--population 0 --episodes 1 --parallel 1\` now prints
"--population must be >= 1, got 0" to stderr and exits 1.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
🧹 Nitpick comments (2)
scripts/benchmarks/bench_evolution_smoke.py (2)
43-45: Adopt NumPy-style docstrings for_parse_argsandmain.Both entry-point functions should use NumPy-style docstrings per repo rules (including a
Returnssection).As per coding guidelines, "Use NumPy-style docstrings".Proposed docstring update
def _parse_args() -> argparse.Namespace: + """Parse command-line arguments for the smoke benchmark. + + Returns + ------- + argparse.Namespace + Parsed CLI arguments. + """ parser = argparse.ArgumentParser(description=__doc__) @@ def main() -> int: - """Run a one-generation benchmark and print wall-clock numbers.""" + """Run a one-generation benchmark and print wall-clock numbers. + + Returns + ------- + int + Exit code. Returns ``0`` on success and ``1`` on user/config errors. + """ args = _parse_args()Also applies to: 67-69
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/benchmarks/bench_evolution_smoke.py` around lines 43 - 45, Add NumPy-style docstrings for the two entry-point functions _parse_args and main: for _parse_args include a short summary, Parameters (if any), and a Returns section specifying that it returns an argparse.Namespace; for main include a short summary, Parameters (if any) and a Returns section (e.g., return code or None) following NumPy docstring formatting conventions; update the docstrings located on the functions named _parse_args and main so they match repository guidelines.
119-130: Tighten theoverridestype annotation to be explicit.
overrides: dictis too broad and weakens static guarantees; use a concrete key/value type.As per coding guidelines, "Use comprehensive type annotations in all code".Proposed typing refinement
- overrides: dict = { + overrides: dict[str, int | bool] = { "generations": 1, "population_size": args.population, "episodes_per_eval": args.episodes, "parallel_workers": args.parallel, "checkpoint_every": 999, # effectively disabled for a 1-gen run }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@scripts/benchmarks/bench_evolution_smoke.py` around lines 119 - 130, The local variable overrides should use a concrete typing instead of the broad dict; change its annotation to a more specific mapping such as Dict[str, Any] (or Dict[str, Union[int, bool]] if you want stricter types) for the overrides dict created before the conditional that checks args.cma_diagonal, and update the import to include typing names if needed; ensure the annotated name (overrides) is used when calling base_evolution.model_copy(update=overrides) so static checkers understand the key/value shapes passed into EvolutionConfig.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@scripts/benchmarks/bench_evolution_smoke.py`:
- Around line 43-45: Add NumPy-style docstrings for the two entry-point
functions _parse_args and main: for _parse_args include a short summary,
Parameters (if any), and a Returns section specifying that it returns an
argparse.Namespace; for main include a short summary, Parameters (if any) and a
Returns section (e.g., return code or None) following NumPy docstring formatting
conventions; update the docstrings located on the functions named _parse_args
and main so they match repository guidelines.
- Around line 119-130: The local variable overrides should use a concrete typing
instead of the broad dict; change its annotation to a more specific mapping such
as Dict[str, Any] (or Dict[str, Union[int, bool]] if you want stricter types)
for the overrides dict created before the conditional that checks
args.cma_diagonal, and update the import to include typing names if needed;
ensure the annotated name (overrides) is used when calling
base_evolution.model_copy(update=overrides) so static checkers understand the
key/value shapes passed into EvolutionConfig.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: afbe9560-0559-43fd-94c8-6143099a842d
📒 Files selected for processing (1)
scripts/benchmarks/bench_evolution_smoke.py
This was referenced Apr 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
cma_diagonal: trueonEvolutionConfig). Dropstell()cost from O(n²) to O(n). At LSTMPPO weight-genome scale (n≈47k),tell()goes from 167 s → 3 ms, making any LSTMPPO+klinotaxis evolution campaign actually feasible.Theme.ASCIIrendering, multi-threaded BLAS, and unconditional f-string construction inlogger.infocalls. All silenced for the evolution worker without changing fitness scores.scripts/benchmarks/bench_evolution_smoke.pywith line-buffered stdout (so partial progress is visible if the run is killed mid-bench).Problem
After M0 (PR #132) shipped, manual LSTMPPO+klinotaxis smoke runs took 2-3 minutes for tiny populations and timed out at 5+ minutes for slightly bigger configs. The M0 task notes had attributed this to "GRU forward pass + 1000-step episode is slow" but that explanation didn't survive profiling.
I profiled a single fitness eval (50 steps): 34 ms total, ~0.7 ms per step. Episode-level work is fine.
I then profiled
CMAESOptimizerat the LSTMPPO genome dim (46,989):That's the actual bottleneck. Standard
cma.CMAEvolutionStrategymaintains a full n×n covariance matrix; at n=47k that's a 17 GB matrix and the cubic-ishtell()update is intractable. The standard mitigation in CMA-ES literature (Ros & Hansen 2008's sep-CMA-ES) is to restrict the covariance to its diagonal — dropstell()to O(n), accepts slower per-generation convergence in exchange for tractability.What changed
Per-step path:
evolution/fitness.py:213— passtheme=Theme.HEADLESStocreate_env_from_config. The fitness eval was defaulting toTheme.ASCII, building and `print()`ing the grid every step in worker processes that have no terminal.evolution/loop.py_init_worker— addtorch.set_num_threads(1)(stops BLAS oversubscription atparallel_workers > 1); silence the sharedquantumnematode.logging_configlogger (whichrunners.pyandagent.pyboth import viafrom quantumnematode.logging_config import logger) to WARNING. Loop's own logger keeps the parent's verbosity so generation-level progress still surfaces.agent/runners.py:767-787,734-737— gate per-steplogger.info/logger.debugwithisEnabledFor(...)so f-strings aren't constructed when the level is filtered.The big one:
optimizers/evolutionary.py— adddiagonal: bool = Falsekeyword arg onCMAESOptimizer.__init__. When True, setsCMA_diagonal=Truein the cma options. Default False preserves back-compat for existing small-genome campaigns.utils/config_loader.py— addcma_diagonal: bool = Falsefield onEvolutionConfigwith a comment explaining when to enable it (genome dim >~1000) and the convergence trade-off.scripts/run_evolution.py— plumbevolution_config.cma_diagonalthrough to the optimiser; updated docstring's stale Timing section to reflect post-perf-fix reality.configs/evolution/lstmppo_foraging_small_klinotaxis.yml— opt in (cma_diagonal: true). The M0 LSTMPPO pilot was previously hours-per-campaign; now sub-minute per generation.configs/evolution/mlpppo_foraging_small.yml— kept atfalsefor byte-identical M0 smoke compatibility, with a comment explaining why and noting real M2 campaigns should opt in.Bench harness:
scripts/benchmarks/bench_evolution_smoke.py— committable timing harness with line-buffered stdout and a--cma-diagonalflag.Skill:
.claude/skills/nematode-run-evolution/skill.md— replaced stale "LSTMPPO is slow per episode" narrative with the real cause (cma_diagonal=Falseon weight genomes); added a genome-dim decision table; updated runtime examples to post-perf-fix numbers.Numbers
LSTMPPO+klinotaxis, population 4, episodes 2, 1 generation, parallel 1:
cma_diagonalcma_diagonal: trueMLPPPO, same params: 0.80 s / 0.100 s per episode.
The MLPPPO genome dim is ~9k — full-cov CMA-ES is borderline at that size. Diagonal mode helps but the per-step fixes also matter. For LSTMPPO at 47k dim, only diagonal mode unblocks the campaign; the per-step fixes alone wouldn't move the needle.
Trade-off honesty
cma_diagonal: truegives up off-diagonal covariance adaptation, so per-generation convergence is slower — typically 2-10× more generations to reach the same fitness on non-separable problems (Ros & Hansen 2008). At large n, full-cov isn't a competing option (you can't run it long enough to find out), so net wall-clock to convergence is dramatically faster with diagonal anyway. The trade-off is documented in the field comment, the parameter docstring, the LSTMPPO pilot config, and the skill.Risks (and how this PR addresses them)
agent.py:958,env.pyrendering methods), not in sensing/RNG/action selection. No fitness drift.quantumnematode.logging_configlogger that all the per-step paths import; the loop's own logger (quantumnematode.evolution.loop) keeps the parent verbosity so progress logs survive.set_num_threads(1)— single-threaded BLAS is bit-identical to multi-threaded for our matmul sizes; no determinism regression.cma_diagonal— defaultFalsepreserves existing M0 campaigns. Opt-in via YAML/Pydantic; existing scenario configs unaffected.Review-pass findings (caught + fixed)
The first review pass caught a real bug:
_init_workerwas silencing logger names that don't actually exist (quantumnematode.agent.runners/quantumnematode.agent.agent). Bothrunners.pyandagent.pyimport their logger viafrom quantumnematode.logging_config import logger, so the actual logger is namedquantumnematode.logging_config. Under--log-level INFO + parallel_workers > 1, the per-step f-string-skip gate would silently fail to fire. Fixed in commitb3771139by silencing the correct logger name. The corresponding test now imports the actual runtime logger and assertsnot isEnabledFor(INFO)— which would have failed under the buggy code.Plus added two new regression tests:
test_cma_diagonal_yaml_propagates_to_optimizer_options— locks the YAML→EvolutionConfig→CMAESOptimizer→cma library options chain end-to-end.test_lstmppo_pilot_config_enables_cma_diagonal— guards against accidentally regressing the LSTMPPO pilot'scma_diagonal: truesetting.A second review pass found nothing else.
Test plan
uv run pytest packages/quantum-nematode/tests/quantumnematode_tests/evolution/ packages/quantum-nematode/tests/quantumnematode_tests/optimizers/test_evolutionary.py -v— all greenuv run pytest -m "not nightly"— 2193 passed in 83 s (was 2191 before this PR; +5 new tests for worker policy, diagonal mode, default-off, pilot-config regression, YAML→optimiser chain)uv run pre-commit run -a— cleanopenspec validate evolution-framework --strict— passes (no spec changes; capability untouched)uv run python scripts/benchmarks/bench_evolution_smoke.py --config configs/evolution/lstmppo_foraging_small_klinotaxis.yml --population 4 --episodes 2locally and confirm sub-minute completioncma_diagonal: false(default) still produces identical fitness scores to before🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Performance
Tests / Chores