Skip to content

✨ Add GNN for RL circuit embedding - #777

Merged
flowerthrower merged 13 commits into
v3/paper-reward-parityfrom
v3/669-gnn-rl
Aug 27, 2026
Merged

flowerthrower merged 13 commits into
v3/paper-reward-parityfrom
v3/669-gnn-rl

Conversation

@flowerthrower

@flowerthrower flowerthrower commented Aug 27, 2026

Copy link
Copy Markdown
Member

🤖 AI text below 🤖

Description

Adds opt-in GNN circuit embeddings to the RL predictor while leaving PredictorEnv and the default flat MaskablePPO path unchanged.

The implementation preserves the prototype's ordered 52-dimensional node features, ordered 36-dimensional global features, GraphSAGE/attention/shared-trunk topology, effective training defaults, separate encoder/head learning rates, and tuned paper configuration.

Graph observations remain exact and variable-sized PyG Data sidecars. A focused ragged rollout buffer batches them with PyG, and a narrow collect_rollouts specialization transports the sidecars through SB3. Masking, callbacks, numeric rollout state, GAE, PPO training/clipping/KL behavior, logging, and .zip serialization remain in MaskablePPO/SB3. There is no graph padding, truncation, or node/edge capacity.

Torch Geometric remains optional through the mqt.predictor[gnn] extra and is imported only when graph mode is enabled. The prototype .pt checkpoint is not migrated because its output layer is coupled to the prototype's 39-action schema. The current stack has 43 actions without the unavailable IBM AI passes, so the preserved architectures contain 331,949 trainable parameters with the generic configuration and 231,142 with GNNConfig.paper().

This raw draft is stacked on #778. New tests, external documentation, and changelog changes are deferred for now, as requested.

Credits @antotu for the prototype implementation in #563 and the experiment branch.

Validation

  • complete repository lint
  • existing flat-policy regression tests (4 passed)
  • masked two-step PPO training and .zip save/load/predict smoke
  • exact 2,100-node / 2,099-edge graph smoke with no padding or capacity limit
  • final stacked parameter-count smoke

Fixes #669

Checklist

  • The pull request only contains commits that are focused and relevant to this change.
  • New tests and documentation are intentionally deferred while the feature stack is assembled.
  • The changes follow the project's style guidelines and introduce no new warnings.
  • The changes are fully tested and pass the CI checks.
  • I have reviewed my own code changes.

If PR contains AI-assisted content:

  • Any agent that created, edited, or submitted GitHub content was explicitly authorized for that scope, as required by our AI Usage Guidelines.
  • Every agent-authored or agent-edited public text body begins with the visible disclosure 🤖 *AI text below* 🤖 (titles are exempt).
  • I have disclosed AI assistance in the PR description.
  • I confirm that I have personally reviewed and understood all AI-generated content, and accept full responsibility for it.

@codecov

codecov Bot commented Aug 27, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 6.13333% with 352 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/mqt/predictor/rl/gnn.py 0.0% 337 Missing ⚠️
src/mqt/predictor/rl/predictor.py 60.5% 15 Missing ⚠️

📢 Thoughts on this report? Let us know!

@flowerthrower
flowerthrower changed the base branch from v3/670-intermediate-rewards to v3/configurable-pass-timeouts August 27, 2026 10:50
flowerthrower and others added 13 commits August 27, 2026 15:05
Assisted-by: GPT-5 via Codex

Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Assisted-by: GPT-5 via Codex

Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Assisted-by: GPT-5.6 via Codex
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Co-authored-by: Antonio Tudisco <anto.tu98@hotmail.it>

Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Signed-off-by: flowerthrower <flowerthrower@users.noreply.github.com>
Assisted-by: GPT 5.6 via Codex
Base automatically changed from v3/configurable-pass-timeouts to v3/paper-reward-parity August 27, 2026 13:09
@flowerthrower
flowerthrower merged commit abeb59c into main Aug 27, 2026
@flowerthrower
flowerthrower deleted the v3/669-gnn-rl branch August 27, 2026 13:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

✨ add GNN for RL circuit embedding

1 participant