feat(rewards): add reward option schemas, a registry, and configuration parsing (#358) - #374
Draft
manzuoni-astera wants to merge 7 commits into
Draft
Conversation
Guidance backprops the value and FK steering picks with argmin, so every reward here is really a loss. Nothing said so. Now the protocol docstring does, and points at the contract test that catches a term with the wrong sign.
The structure-factor reward from diff-use#324 is built in two phases, but nothing in src/ ever called the second one, so it could not run from the pipeline at all. Adds PreparableRewardFunctionProtocol and a prepare_reward_if_needed helper, called from both trajectory scalers once the model atom array exists. prepare() mutates the reward and returns None. The tmol reward in diff-use#319 and the torchref one in diff-use#372 both need this hook. Also replaces an `or` fallback on an AtomArray with a reward_atom_array property. Whether an empty AtomArray is falsy is biotite's call, not ours.
manzuoni-astera
requested review from
a team,
DorisMai,
k-chrispens and
marcuscollins
as code owners
August 13, 2026 21:13
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:13 — with 
GitHub Actions
Error
GitHub Actions
Error
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:13 — with
GitHub Actions
Error
GitHub Actions
Error
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:13 — with
GitHub Actions
Error
GitHub Actions
Error
Contributor
|
Warning Review limit reached
Next review available in: 38 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (19)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR formalizes reward configuration in sampleworks by introducing per-reward option schemas, a central reward registry with lazy builders, and a RewardConfig parser that supports JSON/YAML/TOML. It also adds a two-phase “prepare” hook so rewards that depend on model-specific atom topology (e.g., structure-factor) can be initialized at the correct point in the guidance pipeline.
Changes:
- Added reward option dataclasses and a reward registry (
RewardSpec) that lazily imports builders and supports reward-agnostic experimental-data injection. - Added
RewardConfigfor parsing/validating{reward: {weight, reward_options}}configs (incl. YAML env interpolation) and building single/composite rewards. - Added
PreparableRewardFunctionProtocol+prepare_reward_if_needed()and wired preparation into bothPureGuidanceandFKSteering, using a newreward_atom_arrayselection.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tests/utils/test_guidance_script_utils.py | Updates test to cover reward-agnostic structure loading helper. |
| tests/rewards/test_reward_registry.py | Adds contract tests for registry entries, option coercion, and builder validation. |
| tests/rewards/test_reward_function_contract.py | Clarifies/strengthens reward sign convention documentation in contract tests. |
| tests/rewards/test_reward_config.py | Adds extensive parsing/validation/weight/injection/path-remap tests for RewardConfig. |
| tests/rewards/test_prepare_hook.py | Adds unit tests for the new prepare hook protocol + helper behavior. |
| tests/rewards/test_composite.py | Adds tests for weighted reward combination and build behavior. |
| tests/integration/test_pipeline_integration.py | Verifies preparable rewards are prepared before first evaluation in both trajectory scalers. |
| src/sampleworks/utils/guidance_script_utils.py | Splits structure loading into load_guidance_structure(); keeps deprecated density-only helper via new builder path. |
| src/sampleworks/utils/guidance_constants.py | Extends Rewards enum with STRUCTURE_FACTOR. |
| src/sampleworks/eval/structure_utils.py | Adds reward_atom_array property to standardize reward topology selection. |
| src/sampleworks/core/scalers/pure_guidance.py | Prepares rewards (when needed) after processed structure creation, before denoising loop. |
| src/sampleworks/core/scalers/fk_steering.py | Same prepare hook wiring for FK steering. |
| src/sampleworks/core/rewards/structure_factor.py | Adds registry-compatible builder for structure-factor reward. |
| src/sampleworks/core/rewards/registry.py | Introduces RewardSpec, REWARD_SPECS, option coercion, and single-reward builder entrypoint. |
| src/sampleworks/core/rewards/real_space_density.py | Adds registry-compatible builder for real-space density reward (with missing-input errors). |
| src/sampleworks/core/rewards/protocol.py | Adds minimized-sign convention documentation + preparable reward protocol and helper. |
| src/sampleworks/core/rewards/options.py | Adds option-schema dataclasses and option metadata helpers. |
| src/sampleworks/core/rewards/config.py | Adds RewardConfig parsing/validation, weight resolution, injection, path remapping, and build orchestration. |
| src/sampleworks/core/rewards/composite.py | Adds CompositeReward for weighted sum and preparation forwarding. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+131
to
+135
| hint = typing.get_type_hints(options_cls)[name] | ||
| args = [arg for arg in typing.get_args(hint) if arg is not type(None)] | ||
| if not args: | ||
| return hint | ||
| return args[0] if len(args) == 1 else hint |
Contributor
Author
There was a problem hiding this comment.
Good catch, fixed. option_type now only unwraps unions, so a bare list[str] stays list[str]. No current option declares one, but the next person to write expcolumns: list[str] without | None would have silently lost nargs on its flag. Added a test with a bare generic to hold it.
manzuoni-astera
force-pushed
the
michaelanzuoni/reward-registry-config
branch
from
August 13, 2026 21:33
000b8f5 to
9d03c9f
Compare
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:33 — with
GitHub Actions
Error
GitHub Actions
Error
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:33 — with
GitHub Actions
Error
GitHub Actions
Error
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:33 — with
GitHub Actions
Error
GitHub Actions
Error
The reward was never actually called, so 'no calls before prepare' held trivially. Uses a real DPS step scaler, counts the calls, and drives a mismatch case where the model has four atoms and the structure five, so preparing against the wrong array fails the test. Both from CodeRabbit on diff-use#373.
Reward arguments lived in add_generic_args as if density were the only reward we would ever have, and construction was hardwired to RealSpaceRewardFunction. Each reward now declares its options as a frozen dataclass and registers a lazily-imported builder, so a new reward is a schema plus one registry entry. Builders sit next to their rewards and raise their own missing-input errors, naming both the flag and the config key. CLI flags derive from option names, so flag, config key and schema cannot drift apart. get_reward_function_and_structure splits at its seam: loading the structure is reward-agnostic and becomes load_guidance_structure, the rest is the density builder. Adds Rewards.STRUCTURE_FACTOR, which diff-use#324 never got.
RewardConfig reads the {reward: {weight, reward_options}} mapping from diff-use#358,
as JSON, YAML or TOML. YAML goes through OmegaConf, already a dependency, so
${oc.env:VAR} works the way it does in the run presets.
Weights are 1/N when none are given and verbatim when all are. Giving only
some is an error: a default quietly disagreeing with a number someone typed
is worse than a complaint.
with_experimental_data lets grid search drop in a per-protein map or MTZ
without knowing which reward it is filling.
build_reward turns a configuration into the reward a run scores against. A single reward at full weight comes back as itself, so current runs keep the gradients they have today. Anything else becomes a weighted sum. Weights default to 1/N rather than 1, so adding a term does not quietly scale the gradient up and change what the step size means. Negative weights are rejected: against a minimized objective they flip a term instead of damping it. prepare() forwards to whichever terms need it.
…ard_options option_type stripped None from any hint with type args, so a bare list[str] came back as str and its CLI flag would have lost nargs. Only unions are unwrapped now. reward_options holding a list reached dict() and raised TypeError, which the CLI does not catch, so a typo in a config file printed a traceback. It is a ValueError naming the reward now. Both from Copilot on diff-use#374 and diff-use#375.
manzuoni-astera
force-pushed
the
michaelanzuoni/reward-registry-config
branch
from
August 13, 2026 21:35
9d03c9f to
990351c
Compare
manzuoni-astera
requested a deployment
to
gpu-testing
August 13, 2026 21:35 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
manzuoni-astera
requested a deployment
to
gpu-testing
August 13, 2026 21:35 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
manzuoni-astera
requested a deployment
to
gpu-testing
August 13, 2026 21:35 — with
GitHub Actions
Waiting
GitHub Actions
Waiting
manzuoni-astera
marked this pull request as draft
August 14, 2026 00:01
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
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.
Stacked on #373. GitHub can't take a base branch that only lives in my fork, so this is opened
against
mainand the diff includes #373's commits. Review fromfeat(rewards): add reward option schemas and a reward registryonwards, and merge #373 first.Summary
The data model for #358, with no CLI changes; those are in #375. Reward arguments currently live
in
add_generic_argsas if density were the only reward we will ever have, and rewardconstruction is hardwired to
RealSpaceRewardFunction. Each reward now declares its options onceand registers a builder, so adding a reward type is a schema plus one registry entry rather than
an edit to shared argparse plumbing. Also adds
RewardConfig, which reads the mapping from #358as JSON, YAML or TOML.
Changes
Three modules under
core/rewards/:options.pyholds one frozen dataclass per reward. It is the single declaration that argparse,the config-file schema, container path remapping and the error messages all read, so a flag, its
config key and its schema cannot drift apart. Each option carries metadata for help text,
choices, whether it is a path, and whether it arrives as JSON.
registry.pyholds aRewardSpecper reward: its option schema, a builder addressed as a"module:function"string, which option carries the experimental data, which option carries theresolution, and which options the reward cannot be built without. The string address matters,
because this module has to import in every pixi environment to render
--helpor validate aconfig, while
structure_factor.pyimportsSFC_Torchandreciprocalspaceshipat modulelevel. A reward's real dependencies are only touched when that reward is actually built.
config.pyholdsRewardConfig. Weights resolve to 1/N when none are given and are usedverbatim when all are; giving only some raises, because a default silently disagreeing with a
number someone typed is worse than a complaint.
with_experimental_dataexists for grid search,which resolves a map and a resolution per protein from the CSV long after the reward type was
chosen; each reward declares which of its options those fill, so the injection stays
reward-agnostic. YAML goes through OmegaConf, already a dependency, so
${oc.env:VAR}interpolation behaves the way it does in the run presets.
Builders live next to their rewards and raise their own errors for inputs they cannot do
without, naming both the flag and the config key, which is what Marcus asked for in #319 review.
get_reward_function_and_structuresplits at its seam: loading the structure is reward-agnosticand becomes
load_guidance_structure, the rest is the density builder. Nothing else called it,so it is deleted rather than left as a shim.
Rewards.STRUCTURE_FACTORjoins the enum, which#324 never added.
Testing
tests/rewards/test_reward_registry.pyandtest_reward_config.py, 45 tests. Parsing is checkedagainst the literal YAML from this issue's description, and JSON, YAML and TOML are asserted to
produce equal configurations.
The registry has a contract test that runs over every registered reward: its builder resolves,
its declared data and resolution options name real options, and anything whose name looks like a
file declares
path=True. That last one is a hand-maintained bit of metadata and the easiestthing to forget when adding a reward.
Weight resolution is covered including the partially-weighted error, and the mapping is asserted
to survive both
json.dumpsandpickle, since run configurations go through both.Copilot caught two real defects on this branch, both fixed with tests:
option_typestrippedNonefrom any hint with type arguments, so a barelist[str]collapsed tostrand would havelost
nargson its flag, and areward_optionsvalue that isn't a mapping reacheddict()andraised
TypeErrorinstead of a usage error.CI is green: lint, four typecheck environments, four test environments.
Rollout
Nothing user-facing. No new dependencies, no flags, no configuration anyone has to write, and no
change to how a run is invoked, since the CLI still goes through the old path until #375. Merge
after #373.