feat(cli): select and configure rewards via --reward-type and --reward-config by manzuoni-astera · Pull Request #375 · prism-science/sampleworks

manzuoni-astera · 2026-08-13T21:13:57Z

Stacked on #402 and #374 (#373 has merged), and opened against main for the same reason they
are: a fork branch can't be a base. The diff includes both. Review from feat(cli): select and configure rewards via --reward-type and --reward-config onwards, and merge the other two first.

Summary

Puts the registry from #374 on the command line. --reward-type picks a reward and brings that
reward's own 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, so going from one reward to two is a change to the run, not to the plumbing. The
default 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

sampleworks-guidance ... --reward-type structure_factor --mtzfile 1vme.mtz --bulk-solvent combined
sampleworks-guidance ... --reward-config rewards.yaml

The reward is resolved in the parser's existing first pass, beside --model, because it decides
which flags the second pass registers. Registering only the selected reward's options gets
cross-reward rejection for free: --density under --reward-type structure_factor is an
argparse error, the same way a Boltz flag already is on a Protenix run. Passing a config file and
a --reward-type together is rejected outright rather than one silently winning.

--reward-type defaults to real_space_density, and its flags keep their spellings, so nothing
anyone has typed before changes meaning. tests/cli/test_guidance_cli.py passes untouched, which
is the compatibility check that matters.

Two constraints shaped the serialization, and both are easy to get wrong:

reward_config serializes as a JSON string rather than a nested mapping. as_dict() is written
into the output CIF as the sampleworks category, and add_category_to_cif treats any
non-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_order and em fields stay on GuidanceConfig.
run_grid_search.py builds configs from them, grid_search_eval_utils.py reads density and
resolution back out of job_metadata.json, and job queues are pickled by one build and unpickled
by another. So the two representations are reconciled in one method, in whichever direction has
the information, and __setstate__ fills in reward_config for pickles written before it
existed.

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.

CompositeReward lands here too. build_reward returns a single reward at full weight as
itself, 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-type conflict, the legacy pickle
path, and a metadata round trip that writes the output CIF and reads it back.

tests/rewards/test_reward_build_integration.py runs a command line through to a scored
structure 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, with
RewardConfig.with_experimental_data() as the hook. And #319 conflicts with this by design; once
this is in, it rebases down to a reward class, an options dataclass and one registry entry.

Merge after #373 and #374.

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 prism-science#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 prism-science#319 and the torchref one in prism-science#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.

coderabbitai · 2026-08-13T21:14:11Z

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

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.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands.}

Copilot

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 RewardConfig model 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-phase prepare() 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.

manzuoni-astera · 2026-08-13T21:36:00Z

+            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 {}),
+                )
+            )


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().

…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.

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 prism-science#324 never got.

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.

RewardConfig reads the {reward: {weight, reward_options}} mapping from prism-science#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.

Its two halves are load_guidance_structure and the density builder now, and nothing calls it.

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.

--help read "--mtzfile REWARD_OPTION_MTZFILE". Options with choices keep showing their choices.

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 prism-science#373.

…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 prism-science#374 and prism-science#375.

Only the space-separated spelling was caught.

…inds to Review follow-ups (marcuscollins): - MockPreparableRewardFunction lives in tests/mocks/rewards.py and is used by both the hook unit tests and the trajectory-scaler integration test; the inline RecordingPreparableReward and the private PreparableReward/PlainReward are gone. - test_prepare_hook.py builds atom arrays with build_test_atom_array instead of a one-off helper. - PreparableRewardFunctionProtocol docstring says which model: the generative model's atom array (reward_atom_array), not the processed input structure. - Both scaler call sites explain that to_reward_inputs does not mutate the frozen processed structure or reconcile coordinates, why reward_atom_array is the right topology for prepare, and where the reconciled reference coords live.

…373) ## Summary The structure-factor reward merged in #324 cannot run. It is built in two phases, and nothing in `src/` has ever called the second one, so the reward has tests but no way to reach a guidance job. This adds a prepare hook to the reward protocol and calls it from both trajectory scalers once the model atom array exists. It makes #324's reward reachable and gives #319 and #372 one hook to share instead of three private variants. Independent of the registry work in #374 and #375. ## Changes `PreparableRewardFunctionProtocol` joins `core/rewards/protocol.py`, alongside a `prepare_reward_if_needed(reward, atom_array, *, device)` helper that no-ops on rewards which don't implement it. Callers apply it unconditionally, so the density path is untouched. The signature is the one #324 already shipped: `prepare(atom_array, *, device) -> None`, mutating the reward rather than returning a new object. Worth settling now, because the two rewards queued behind this disagree with it. #319's prepare returns a value, and #372's takes a structure rather than an atom array. Both call sites sit in `sample()`, right after the processed structure is built and before the denoising loop: `core/scalers/pure_guidance.py` and `core/scalers/fk_steering.py`. They pass `processed_structure.reward_atom_array`, a new property that returns the model atom array when the model exposes one and the input structure's otherwise. That is deliberately the same choice `to_reward_inputs` makes. If the two ever diverge, a reward's atom ordering stops matching the coordinate tensor it is scoring, and the failure is silent: wrong numbers, no error. The property also replaces `model_atom_array or atom_array`. Whether an empty `AtomArray` is falsy is biotite's decision, not something we should build on. Also in here: the protocol docstring now states that rewards are minimized, since guidance backprops the value and FK steering selects with `argmin`. Nothing said so before, and weighted combinations in #375 only mean something if every term agrees on the sign. ## Testing `tests/rewards/test_prepare_hook.py` covers the helper directly: it prepares a two-phase reward with the atom array and device it was given, leaves a one-phase reward alone, and can be re-run to rebind a reward to a different topology. `tests/integration/test_pipeline_integration.py` drives both trajectory scalers end to end with a recording reward, asserting that `prepare` ran exactly once, before the first evaluation, with the model's atom array. The model there has four atoms and the structure five, so preparing against the input structure fails the test rather than passing by coincidence. A real `DataSpaceDPSScaler` is used so the reward is genuinely evaluated. Both of those came out of CodeRabbit review on this PR, and the first version of the test passed vacuously without them. CI is green: lint, four typecheck environments, four test environments. The GPU workflow has not run. ## Rollout Nothing. No new dependencies, no CLI or configuration changes, and no behavior change for any reward that doesn't implement `prepare`. Merge this before #374 and #375, which build on it. --------- Co-authored-by: xraymemory <me.anzuoni@gmail.com> Co-authored-by: M E A <xraymemory@users.noreply.github.com>

Decided in the 4 Sep rewards discussion. `to_reward_inputs` reconciles the structure's B-factors and reference coordinates onto the model atom order, but those values live only in the returned RewardInputs; the model atom array keeps its template placeholders and nothing writes back into it, on purpose. Preparing the SF reward from that array built the SFcalculator, and its solvent estimate, from placeholder coordinates. `PreparableRewardFunctionProtocol.prepare(reward_inputs, *, device)` and `prepare_reward_if_needed` now take the RewardInputs the scalers attach to the step context, so a reward is prepared against exactly what its __call__ is fed. RewardInputs gains `atom_array` (the topology it was built from, set by from_atom_array) and `to_atom_array()`, which copies that topology with coord/b_factor from the tensors and occupancy 1.0. StructureFactorRewardFunction.prepare builds its gemmi structure from it. Tests: prepare hook receives the inputs object itself; to_atom_array round trip, reconciled override, template mismatch, no-topology error; mismatch integration asserts model topology + reconciled values and that the model atom array is untouched; SF fixtures wrap their atom array in RewardInputs.

…/reward-registry-config

…/reward-cli

…ism-science#373 Forward the reward inputs to each component instead of an atom array, and build RewardInputs in the tests that prepare a reward by hand.

…dance_structure Merging main brought two tests written against get_reward_function_and_structure; on this branch the temporary altloc CIF is handled by load_guidance_structure, so they exercise that (as the earlier keeps-original-file test already does).

…oni/reward-cli

Decided in the 4 Sep rewards discussion. `to_reward_inputs` reconciles the structure's B-factors and reference coordinates onto the model atom order, but those values live only in the returned RewardInputs; the model atom array keeps its template placeholders and nothing writes back into it, on purpose. Preparing the SF reward from that array built the SFcalculator, and its solvent estimate, from placeholder coordinates. `PreparableRewardFunctionProtocol.prepare(reward_inputs, *, device)` and `prepare_reward_if_needed` now take the RewardInputs the scalers attach to the step context, so a reward is prepared against exactly what its __call__ is fed. RewardInputs gains `atom_array` (the topology it was built from, set by from_atom_array) and `to_atom_array()`, which copies that topology with coord/b_factor from the tensors and occupancy 1.0. StructureFactorRewardFunction.prepare builds its gemmi structure from it. Tests: prepare hook receives the inputs object itself; to_atom_array round trip, reconciled override, template mismatch, no-topology error; mismatch integration asserts model topology + reconciled values and that the model atom array is untouched; SF fixtures wrap their atom array in RewardInputs.

…lanzuoni/reward-registry-config

…oni/reward-cli

xraymemory added 2 commits August 13, 2026 16:45

docs(rewards): write down that rewards are minimized

c9aa02a

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.

Copilot AI lite review requested due to automatic review settings August 13, 2026 21:13

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

manzuoni-astera mentioned this pull request Aug 13, 2026

feat(rewards): add a prepare hook to the reward protocol and scalers #373

Merged

Copilot started reviewing on behalf of manzuoni-astera August 13, 2026 21:14 View session

Copilot AI reviewed Aug 13, 2026

View reviewed changes

manzuoni-astera force-pushed the michaelanzuoni/reward-cli branch from 0eae3c9 to 6994697 Compare August 13, 2026 21:33

manzuoni-astera had a problem deploying to gpu-testing August 13, 2026 21:33 — with GitHub Actions Error

xraymemory added 9 commits August 13, 2026 17:34

docs(rewards): document reward selection and how to add a reward type

9520db2

--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.

test(rewards): check a reward is usable end to end from the command line

b6d5d6d

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.

refactor(rewards): drop get_reward_function_and_structure

782ada3

Its two halves are load_guidance_structure and the density builder now, and nothing calls it.

fix(cli): show reward option metavars instead of their internal dest

a27e85f

--help read "--mtzfile REWARD_OPTION_MTZFILE". Options with choices keep showing their choices.

xraymemory added 2 commits August 13, 2026 17:34

fix(cli): reject --reward-type=value beside --reward-config

6683c29

Only the space-separated spelling was caught.

manzuoni-astera force-pushed the michaelanzuoni/reward-cli branch from 6994697 to 6683c29 Compare August 13, 2026 21:35

manzuoni-astera had a problem deploying to gpu-testing August 13, 2026 21:35 — with GitHub Actions Error

manzuoni-astera marked this pull request as draft August 14, 2026 00:01

This was referenced Aug 14, 2026

Build a reward configuration system #358

Open

feat(rewards): add reward option schemas, a registry, and configuration parsing (#358) #374

Draft

xraymemory and others added 2 commits August 28, 2026 11:32

Merge branch 'main' into michaelanzuoni/reward-prepare-hook

f2bead8

xraymemory added 6 commits September 4, 2026 23:10

Merge branch 'michaelanzuoni/reward-prepare-hook' into michaelanzuoni…

51f9594

…/reward-registry-config

Merge branch 'michaelanzuoni/reward-prepare-hook' into michaelanzuoni…

2f30e89

…/reward-cli

fix(rewards): CompositeReward.prepare takes RewardInputs, matching pr…

6b14beb

…ism-science#373 Forward the reward inputs to each component instead of an atom array, and build RewardInputs in the tests that prepare a reward by hand.

Merge branch 'michaelanzuoni/reward-registry-config' into michaelanzu…

dfad17b

…oni/reward-cli

manzuoni-astera had a problem deploying to gpu-testing September 5, 2026 03:46 — with GitHub Actions Error

xraymemory mentioned this pull request Sep 5, 2026

feat(rewards): prepare() takes RewardInputs, not an atom array #402

Open

xraymemory added 2 commits September 5, 2026 00:44

Merge branch 'michaelanzuoni/prepare-takes-reward-inputs' into michae…

f578a9f

…lanzuoni/reward-registry-config

Merge branch 'michaelanzuoni/reward-registry-config' into michaelanzu…

98282e5

…oni/reward-cli

manzuoni-astera requested a deployment to gpu-testing September 5, 2026 04:45 — with GitHub Actions Waiting

Navigation Menu

Uh oh!

feat(cli): select and configure rewards via --reward-type and --reward-config - #375

feat(cli): select and configure rewards via --reward-type and --reward-config#375
manzuoni-astera wants to merge 24 commits into
prism-science:mainfrom
manzuoni-astera:michaelanzuoni/reward-cli

manzuoni-astera commented Aug 13, 2026 •

edited by xraymemory

Loading

Uh oh!

coderabbitai Bot commented Aug 13, 2026 •

edited

Loading

Draft PR not reviewed

Uh oh!

Copilot AI left a comment

Uh oh!

manzuoni-astera Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

manzuoni-astera commented Aug 13, 2026 • edited by xraymemory Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Testing

Rollout

Uh oh!

coderabbitai Bot commented Aug 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Draft PR not reviewed

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

manzuoni-astera Aug 13, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

manzuoni-astera commented Aug 13, 2026 •

edited by xraymemory

Loading

coderabbitai Bot commented Aug 13, 2026 •

edited

Loading