feat: Transformer (temporal-window) PPO brain â Gate-2 plugin parity (§5) - #212
Merged
Conversation
âĶiew) Pre-implementation readiness review of §5 (Gate-2 transformer parity). Verified the plan against current code (âĪ6-file budget holds â config map auto-derived but BrainConfigType union + 5 other sites stand; CLI-default match block still needs the 5.3 registry refactor; 24 brains â 24â25 correct). Settled the two open design choices: - Transformer input = temporal window (rolling buffer of last N feature vectors â encoder), an attention-based temporal-memory comparator alongside LSTM/CfC. - Both action heads: discrete (G2.b/G2.c parity) + continuous tanh-Gaussian via the shared _policy (no extra files; continuous arm joins the §6 tuning). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
âĶrch branch) Replace the `match brain_type` CLI-default block in run_simulation.py with a registry-driven default: `get_registration(brain_type.value).config_cls(seed=âĶ)`. Removes the per-architecture branch at the entrypoint (the G2.c "no per-arch branches" claim must hold here too), so a new architecture â e.g. the §5 transformer â needs no `case` here. One-time platform fix, landed as a SEPARATE prior commit so it is NOT counted in the transformer-addition `git diff --name-only` that measures the âĪ6-file plugin-parity budget (G2.b). §5.3. Behaviour preserved: `brain_type` defaults to `DEFAULT_BRAIN_TYPE` (qvarcircuit, seed-only-constructible) and is overwritten from the config file below, so config runs are unaffected (verified). Brains whose config requires fields (e.g. `sensory_modules`) must be run via `--config` â the seed-only default now raises a clear pydantic ValidationError instead of the previous stale silent fall-through (the match listed `mlpppo`, which can't construct seed-only, and omitted `connectomeppo`, which can). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
âĶarity (§5.1/§5.2) Add a Transformer self-attention PPO brain as the Gate-2 plugin-parity vehicle (G2.b/G2.c). The policy/value trunk is a Transformer encoder over a temporal window of the last N sensory feature vectors â an attention-based temporal-memory comparator alongside LSTM/CfC. Both action heads: discrete (Categorical) and continuous (tanh-squashed Gaussian) via the shared `_policy` helpers. Window-as-state: each step stores its own front-zero-padded (window_size, input_dim) window in the shared RolloutBuffer, so each entry is self-contained and the update re-forwards the stored window â reusing the shuffled-minibatch buffer (the window *is* the temporal context; no recurrent BPTT chain). Plugin-parity (G2.b): the addition touches exactly 5 files (âĪ 6) â new module, `dtypes.py` enum + Literal, `brain/arch/__init__.py` import/export, `config_loader.py` BrainConfigType union, and the test. NO `_build_infra_kwargs` branch (the brain matches the default infra shape) and NO per-architecture branch anywhere (registry-driven CLI default landed separately in the prior commit) â G2.c. The YAML config example + docs enumerations land in separate commits, outside the measured diff. Tests (8): registration, discrete/continuous head shapes, front-padded window + per-episode reset, discrete-grid + continuous-2D training smokes (no NaN, worm moves), weight round-trip. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Outside the G2.b parity diff (separate from the measured transformer-addition commit): - Ship `transformerppo_small_klinotaxis.yml` (discrete grid) â the §5.5 smoke. 200-run headless training shows a clean learning signal: quarter-mean return 5.2 â 36.2 (last > first) and foods/episode 3.06 â 9.94 (â target 10; 196/200 runs ate âĨ1), no NaN/Inf. - Bump the brain-arch enumerations 24 â 25 + add `transformerppo`: AGENTS.md, README.md, CONTRIBUTING.md, openspec/config.yaml (§5.6 â platform bookkeeping, not per-architecture parity cost). - Tick §5.1â5.6 with outcomes. G2.b measured = 5 files (âĪ6); G2.c = no per-architecture branches added; the GO/PIVOT/STOP verdict lands in §7. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Cosmetic only â ruff-format wasn't run on the new file before the commit-B snapshot. No logic change; G2.b's 5-file measurement (on the isolated addition commit) is unaffected (this touches an already-counted file). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
D8 captured the parity mechanics but was silent on the realized implementation. Add an implementation note: G2.b realized = 5 files (no _build_infra_kwargs branch, no BRAIN_CONFIG_MAP edit â auto-derived); architecture = temporal-window self-attention with window-as-state reusing the shared RolloutBuffer (no BPTT buffer); 5.3 also fixed a stale CLI match (mlpppo/connectomeppo); transformer learns cleanly on klinotaxis (bonus). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
âĶiew)
Branch review: an invalid d_model/nhead combination raised a raw torch
AssertionError ('embed_dim must be divisible by num_heads') deep in brain
construction. Add a pydantic model_validator (matching the other brain configs)
so it fails fast with a clear message at config time; also guards window_size>=1.
Test added.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. ð âđïļ Recent review infoâïļ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: ð Files selected for processing (2)
ð§ Files skipped from review as they are similar to previous changes (2)
ð WalkthroughWalkthroughAdds a TransformerPPO brain: a Transformer-encoder PPO agent over rolling temporal windows, integrated into enums, config loading, registry-driven runtime selection, tests, an example foraging scenario, and updated docs/specs enumerating 25 brain architectures. ChangesTransformerPPO Brain Architecture
Estimated code review effortðŊ 4 (Complex) | âąïļ ~45 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
June 6, 2026 14:33
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
ðĪ Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@openspec/changes/add-continuous-2d-and-action-heads/design.md`:
- Line 89: The implementation note header currently reads "Implementation note
(2026-06-07):" which is a future date; update that header in the document (the
line containing "Implementation note (2026-06-07):") to either the actual
completion date "2026-06-06" or remove the parenthetical date entirely so the
record is not out-of-order.
In `@packages/quantum-nematode/quantumnematode/brain/arch/transformer_ppo.py`:
- Around line 139-160: TransformerPPOBrain currently accepts a mutable
DEFAULT_ACTIONS list by reference (action_set parameter, self._action_set, and
the action_set getter/setter), which lets callers mutate a shared list across
instances; fix this by changing the action_set default to None (or accept an
iterable) and ensure the constructor makes a defensive copy and stores an
immutable sequence (e.g., tuple(self._action_set) or
list(self._action_set).copy()), and update the action_set property to return an
immutable copy and to store a defensive copy on set so mutations cannot leak
between instances; refer to TransformerPPOBrain, the action_set parameter,
self._action_set, and the action_set getter/setter when making the changes.
In
`@packages/quantum-nematode/tests/quantumnematode_tests/brain/arch/test_transformer_ppo.py`:
- Around line 23-37: Add a NumPy-style docstring to the _config function:
describe the purpose of the helper, document the parameter continuous (bool)
under Parameters with its meaning, and document the return value as a
TransformerPPOBrainConfig under Returns; keep concise and follow NumPy-style
sections (Parameters, Returns) and include types and a one-line description;
place this docstring directly above the function definition for _config in the
test file so tools and linters recognize it.
ðŠ 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: c71c8929-e25f-49f4-91ce-3b8296cbfa4e
ð Files selected for processing (13)
AGENTS.mdCONTRIBUTING.mdREADME.mdconfigs/scenarios/foraging/transformerppo_small_klinotaxis.ymlopenspec/changes/add-continuous-2d-and-action-heads/design.mdopenspec/changes/add-continuous-2d-and-action-heads/tasks.mdopenspec/config.yamlpackages/quantum-nematode/quantumnematode/brain/arch/__init__.pypackages/quantum-nematode/quantumnematode/brain/arch/dtypes.pypackages/quantum-nematode/quantumnematode/brain/arch/transformer_ppo.pypackages/quantum-nematode/quantumnematode/utils/config_loader.pypackages/quantum-nematode/tests/quantumnematode_tests/brain/arch/test_transformer_ppo.pyscripts/run_simulation.py
|
|
||
| Add a Transformer/attention PPO brain (`transformer-brain`) through the plugin interface as the **G2.b/G2.c verification** (files-touched âĪ 6; no per-arch branches), recording engineer-hours for G2.a. It doubles as the roadmap's MAY transformer comparison row. To make the no-per-arch-branch claim hold at the entrypoint too, refactor the `match brain_type` CLI-default block ([`scripts/run_simulation.py:286-298`](../../../scripts/run_simulation.py#L286)) into a registry-driven default. **This refactor ships in a SEPARATE prior commit** â it is a one-time platform fix, not a per-architecture cost â so it is NOT counted in the transformer-addition `git diff --name-only` that measures G2.b. **Files-budget measurement (reconciling the two enumerations):** G2.b's threshold is `âĪ 6 files`; its parenthetical list in [`phase6-tracking` design.md Decision 6](../phase6-tracking/design.md) ("registry registration + brain implementation + config class + config example + smoke test + docs") is illustrative. The *operative* enumeration is the [plugin-developer-guide](../../../docs/architecture/plugin-developer-guide.md) budget: new module (brain + config class + `@register_brain`), `dtypes.py` enum member, `brain/arch/__init__.py` import/export, `config_loader.py` `BrainConfigType` union, `brain_factory.py` `_build_infra_kwargs` branch *only if non-default*, and a test file = âĪ 6, with the YAML config example explicitly *outside* the 6 â as are the cross-repo brain-arch enumeration updates (`AGENTS.md`, `README.md`, `CONTRIBUTING.md`, `openspec/config.yaml`, currently "24 architectures"), which are required for merge but are platform-bookkeeping, not per-architecture parity cost (landed in a separate commit from the runtime wiring so the G2.b `git diff` stays clean). The âĪ 6 threshold itself is unchanged (this is a pre-fire measurement clarification, not a gate recalibration). **Why a real arch over a hypothetical baseline:** the user chose the transformer to get a reusable comparator; G2.a explicitly makes engineer-hours non-load-bearing precisely so a heavier arch doesn't distort the parity verdict. **Gate-2 floor check (G2.d):** connectome + MLP-PPO train cleanly on continuous-2D klinotaxis to âĨ 50 % of their T4 grid-substrate mean episode return (the 50 % floor is wide because discrete-grid and continuous-2D return scales aren't directly comparable; this checks "didn't break", not ranking). | ||
|
|
||
| **Implementation note (2026-06-07):** |
There was a problem hiding this comment.
Fix future-dated implementation note timestamp.
Line 89 says 2026-06-07, which is ahead of the current timeline (2026-06-06) and makes the record look out-of-order. Please use the actual completion date or remove the date until that day passes.
ðĪ Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@openspec/changes/add-continuous-2d-and-action-heads/design.md` at line 89,
The implementation note header currently reads "Implementation note
(2026-06-07):" which is a future date; update that header in the document (the
line containing "Implementation note (2026-06-07):") to either the actual
completion date "2026-06-06" or remove the parenthetical date entirely so the
record is not out-of-order.
Codecov Reportâ Patch coverage is
ðĒ Thoughts on this report? Let us know! |
Address PR #212 review: - TransformerPPOBrain now defensively copies action_set on store (__init__ + setter): `self._action_set = list(action_set)`. Prevents the shared module-level DEFAULT_ACTIONS default arg from being mutated across instances â matches the CfC/connectome precedent. (Signature kept = DEFAULT_ACTIONS for consistency with all sibling brains; the store-copy fully addresses the leak.) - Add a NumPy-style docstring to the `_config` test helper. Skipped: the design.md "Implementation note (2026-06-07)" date â verified `date` is actually 2026-06-07 (the session date rolled mid-work), so it is the real authoring date and chronological with the 2026-06-06 §4 notes, not a future or out-of-order date. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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
§5 of
add-continuous-2d-and-action-heads: add a Transformer self-attention PPO brain as the Gate-2 plugin-parity vehicle (G2.b/G2.c) and the roadmap's MAY transformer comparison row. The trunk is a Transformer encoder over a temporal window of recent sensory feature vectors â an attention-based temporal-memory comparator alongside LSTM/CfC. Both action heads (discrete + continuous via the shared_policy).Commit structure (so the G2.b parity diff is clean)
refactor(run-sim)â registry-driven CLI-default brain config (removes the per-archmatch brain_typebranch at the entrypoint). One-time platform fix, separate prior commit, not counted in G2.b.feat(transformer)â the measured transformer-addition: 5 files (âĪ 6).docs(transformer)â klinotaxis config + 24â25 enumerations + task ticks (outside the parity diff).Gate-2 results
git show --name-onlyon commit 2): new module,dtypes.pyenum,brain/arch/__init__.py,config_loader.pyBrainConfigTypeunion, test. No_build_infra_kwargsbranch (default infra shape);BRAIN_CONFIG_MAPauto-derives from the registry.isinstance(brain, âĶ)added.Design choices
(window_size, input_dim)window in the shared shuffled-minibatchRolloutBuffer; the update re-forwards the stored window. Self-contained entries â no recurrent BPTT buffer despite being temporal._policy; learnable positional embedding; shared encoder + separate actor/critic heads.matchlistedmlpppo, which can't construct seed-only, and omittedconnectomeppo).Validation
OpenSpec
§5.1â5.6 ticked; D8 carries the implementation-outcome note;
openspec validate --strictpasses.ðĪ Generated with Claude Code
Summary by CodeRabbit
New Features
Documentation
Tests