feat(cli): select and configure rewards via --reward-type and --reward-config - #375
Draft
manzuoni-astera wants to merge 13 commits into
Draft
feat(cli): select and configure rewards via --reward-type and --reward-config#375manzuoni-astera wants to merge 13 commits into
manzuoni-astera wants to merge 13 commits into
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,
Abdelsalam-Abbas,
DorisMai,
denis-zaitsev,
k-chrispens,
mag-astera and
marcuscollins
as code owners
August 13, 2026 21:13
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:14 — with 
GitHub Actions
Error
GitHub Actions
Error
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:14 — with
GitHub Actions
Error
GitHub Actions
Error
manzuoni-astera
had a problem deploying
to
gpu-testing
August 13, 2026 21:14 — 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 (25)
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 makes rewards a first-class, selectable part of sampleworks-guidance: runs can choose a single reward via --reward-type (auto-registering only that reward’s flags) or define one-or-more rewards in a --reward-config file that parses into a unified RewardConfig. It also adds a “prepare” hook so topology-dependent rewards (e.g., structure factors) can bind to the model’s atom ordering before the first evaluation, and introduces CompositeReward for weighted combinations.
Changes:
- Added a reward registry + per-reward option schemas, and a
RewardConfigmodel that parses JSON/YAML/TOML and can be serialized safely into run metadata. - Updated the CLI parsing flow to resolve reward selection early, generating only the selected reward’s option flags and rejecting cross-reward flags automatically.
- Implemented composable multi-reward guidance (
CompositeReward) and a two-phaseprepare()hook invoked by trajectory scalers.
Reviewed changes
Copilot reviewed 25 out of 25 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 tests to use load_guidance_structure() after reward building is decoupled. |
| tests/utils/test_guidance_script_arguments.py | Adds tests for GuidanceConfig reward-config reconciliation, legacy pickle migration, and safe metadata serialization. |
| tests/rewards/test_reward_registry.py | Contract tests for the reward registry and option schema coercion. |
| tests/rewards/test_reward_function_contract.py | Documents/enforces the “rewards are minimized” sign convention in contract tests. |
| tests/rewards/test_reward_config.py | Adds parsing/validation/weighting tests for reward config files across formats. |
| tests/rewards/test_reward_build_integration.py | GPU-marked end-to-end tests from CLI argv → built reward → prepare() → scoring. |
| tests/rewards/test_prepare_hook.py | Unit tests for the PreparableRewardFunctionProtocol and helper. |
| tests/rewards/test_composite.py | Tests weighted reward composition, gradients, validation, and builder integration. |
| tests/integration/test_pipeline_integration.py | Ensures preparable rewards are prepared before first call in both trajectory scalers. |
| tests/cli/test_guidance_cli.py | Adds CLI coverage for reward selection, config-file composition, and flag rejection behavior. |
| src/sampleworks/utils/guidance_script_utils.py | Splits structure loading from reward building; builds rewards via RewardConfig + registry. |
| src/sampleworks/utils/guidance_script_arguments.py | Adds --reward-type / --reward-config, generates per-reward flags from schemas, and reconciles legacy density fields. |
| src/sampleworks/utils/guidance_constants.py | Extends Rewards enum with STRUCTURE_FACTOR. |
| src/sampleworks/eval/structure_utils.py | Introduces reward_atom_array to consistently choose model-vs-structure atom topology for rewards. |
| src/sampleworks/core/scalers/pure_guidance.py | Calls prepare_reward_if_needed() once topology is known. |
| src/sampleworks/core/scalers/fk_steering.py | Calls prepare_reward_if_needed() once topology is known. |
| src/sampleworks/core/rewards/structure_factor.py | Adds a registry builder entrypoint for structure-factor reward construction. |
| src/sampleworks/core/rewards/registry.py | Implements reward registry, lazy builder import, and option coercion. |
| src/sampleworks/core/rewards/real_space_density.py | Adds a registry builder entrypoint for density reward construction. |
| src/sampleworks/core/rewards/protocol.py | Adds sign-convention docs and the preparable reward protocol + helper. |
| src/sampleworks/core/rewards/options.py | Defines frozen dataclass option schemas used by CLI/config/metadata. |
| src/sampleworks/core/rewards/config.py | Implements RewardConfig parsing/validation, defaults materialization, path remapping, and reward building. |
| src/sampleworks/core/rewards/composite.py | Implements CompositeReward for weighted combinations and preparation forwarding. |
| README.md | Documents reward selection and config-file composition in user-facing docs. |
| AGENTS.md | Updates agent guidance for reward selection, config files, and how to add new reward types. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+185
to
+192
| weight = entry.get(WEIGHT_KEY) | ||
| entries.append( | ||
| RewardEntry( | ||
| reward=reward, | ||
| weight=None if weight is None else float(weight), | ||
| options=dict(entry.get(REWARD_OPTIONS_KEY) or {}), | ||
| ) | ||
| ) |
Contributor
Author
There was a problem hiding this comment.
Fixed. reward_options is checked for mapping-ness and raises a ValueError naming the reward and the expected shape, so a typo in a config file reads as a usage error instead of a TypeError traceback out of dict().
manzuoni-astera
force-pushed
the
michaelanzuoni/reward-cli
branch
from
August 13, 2026 21:33
0eae3c9 to
6994697
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.
…d-config --reward-type picks a reward and brings that reward's flags with it, generated from its schema. --reward-config takes the same configuration from a file and is the only way to combine rewards. Both produce one RewardConfig. The reward is resolved in the existing first parse pass, beside --model, because it decides which flags exist. Registering only the selected reward's options gets us cross-reward rejection for free, the way a Boltz flag is already rejected on a Protenix run. The default is real_space_density and its flags keep their spellings, so existing command lines, presets and CLI tests are untouched. GuidanceConfig keeps the flat density fields in step with the configuration both ways: grid search and old pickles build from the flat fields, the eval scripts read density and resolution back out of job_metadata.json. reward_config serializes as a JSON string. as_dict() also becomes a CIF category, and add_category_to_cif reads any non-string iterable as a column of rows. A missing required input is a usage error, raised before a model loads.
--density is no longer a fact about every run, so the README and AGENTS.md now say which options belong to which reward. Structure factors move to implemented in the data-types list.
Covers argv, configuration, build, prepare, score. Every piece of that had tests; the seams between them did not, which is how the structure-factor reward merged without being runnable.
Its two halves are load_guidance_structure and the density builder now, and nothing calls it.
--help read "--mtzfile REWARD_OPTION_MTZFILE". Options with choices keep showing their choices.
Only the space-separated spelling was caught.
manzuoni-astera
force-pushed
the
michaelanzuoni/reward-cli
branch
from
August 13, 2026 21:35
6994697 to
6683c29
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 was referenced Aug 14, 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
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 and #374, and opened against
mainfor the same reason they are: a fork branchcan't be a base. The diff includes both. Review from
feat(cli): select and configure rewards via --reward-type and --reward-configonwards, and merge the other two first.Summary
Puts the registry from #374 on the command line.
--reward-typepicks a reward and brings thatreward's own flags with it, generated from its schema;
--reward-configtakes the sameconfiguration from a file and is the only way to combine rewards. Both produce one
RewardConfig, so going from one reward to two is a change to the run, not to the plumbing. Thedefault is unchanged, so existing command lines, presets and tests keep working. This is also
what finally makes the structure-factor reward from #324 runnable.
Changes
The reward is resolved in the parser's existing first pass, beside
--model, because it decideswhich flags the second pass registers. Registering only the selected reward's options gets
cross-reward rejection for free:
--densityunder--reward-type structure_factoris anargparse error, the same way a Boltz flag already is on a Protenix run. Passing a config file and
a
--reward-typetogether is rejected outright rather than one silently winning.--reward-typedefaults toreal_space_density, and its flags keep their spellings, so nothinganyone has typed before changes meaning.
tests/cli/test_guidance_cli.pypasses untouched, whichis the compatibility check that matters.
Two constraints shaped the serialization, and both are easy to get wrong:
reward_configserializes as a JSON string rather than a nested mapping.as_dict()is writteninto the output CIF as the
sampleworkscategory, andadd_category_to_ciftreats anynon-string iterable as a column of rows, so a nested dict there gets iterated into its keys and
produces a ragged category at the very end of an expensive run.
The flat
density,resolution,loss_orderandemfields stay onGuidanceConfig.run_grid_search.pybuilds configs from them,grid_search_eval_utils.pyreads density andresolution back out of
job_metadata.json, and job queues are pickled by one build and unpickledby another. So the two representations are reconciled in one method, in whichever direction has
the information, and
__setstate__fills inreward_configfor pickles written before itexisted.
A reward missing a required input is reported as a usage error before a model is loaded. The
builders check the same thing for callers that arrive another way.
CompositeRewardlands here too.build_rewardreturns a single reward at full weight asitself, so current runs keep the gradients they have, and anything else becomes a weighted sum
with 1/N defaults. README and AGENTS.md gain a section on picking rewards and on what adding a
reward type involves.
Testing
Full fast suite green, 972 tests. New CLI coverage for reward selection, config files in three
formats, cross-reward flag rejection, the config-plus-
--reward-typeconflict, the legacy picklepath, and a metadata round trip that writes the output CIF and reads it back.
tests/rewards/test_reward_build_integration.pyruns a command line through to a scoredstructure for both rewards and for a two-reward config. That seam is what let #324 merge
unreachable, so it now has a test of its own.
CI is green: lint, four typecheck environments, four test environments.
Not covered: the GPU workflow has not run on any of the stack, so the structure-factor numerics
have only been exercised on CPU. The gpu-marked reward tests are 122 passed and 21 failed on my
laptop, with every failure being "Torch not compiled with CUDA enabled", and the same tests fail
identically on
main. Someone with a GPU box should confirm before merge.Rollout
No migration. Existing command lines, preset TOMLs, pickled job queues and the evaluation scripts
all keep working, and the new flags are optional. Docs ship in this PR.
Two things to know after it lands. Grid search is still density-only: the proteins CSV is
structure,density,resolution, and structure factors would need it to carry an MTZ, withRewardConfig.with_experimental_data()as the hook. And #319 conflicts with this by design; oncethis is in, it rebases down to a reward class, an options dataclass and one registry entry.
Merge after #373 and #374.