[reward, rollout] fix: align sparse reward_extra_info across batch paths by keepkeen · Pull Request #7151 · verl-project/verl

keepkeen · 2026-07-26T10:47:09Z

What does this PR do?

Fixes sparse reward_extra_info assembly in both framework-owned batch paths:

RewardLoopManager.compute_rm_score
AgentLoopWorker._postprocess

Both paths inferred the batch schema from sample 0 and then indexed every sample with those keys. If later samples omitted a key, assembly raised KeyError; if sample 0 had fewer keys, later fields silently disappeared from non_tensor_batch and reward_extra_keys.

This is reachable with verl's built-in scorers, not only malformed custom callbacks: default_compute_score("math_dapo", ...) returns {"score", "acc", "pred"}, while default_compute_score("openai/gsm8k", ...) returns a scalar that the reward manager records as {"acc"}. Their ordering therefore determined whether a mixed-schema batch crashed or lost columns.

The fix introduces a small NumPy-only shared assembler that:

takes the ordered union of keys across all samples;
preserves first-seen key order;
fills absent values with None using dtype=object only for sparse columns;
preserves NumPy's inferred dtype for fully populated columns;
is used by both reward-loop and agent-loop assembly.

Agent-loop workers keep the raw per-sample dictionaries in a stable object column and
the manager assembles them only after all worker chunks are concatenated. The same
finalizer is used by the fully-async single-worker return path. This avoids
reintroducing the failure when different schemas land on different workers or when
the normal manager concat path is bypassed.
Collisions between flattened reward keys and existing batch fields are rejected with
a descriptive error instead of silently overwriting data.

For example:

assemble_reward_extra_info([
    {"acc": 1.0},
    {"score": 1.0, "acc": True, "pred": "42"},
])
# {
#   "acc": array([1.0, 1.0]),
#   "score": array([None, 1.0], dtype=object),
#   "pred": array([None, "42"], dtype=object),
# }

Relationship to earlier work

[reward] fix: union reward_extra_info keys across samples in compute_rm_score #6545 proposed union-key assembly only in RewardLoopManager and was closed with the expectation that users provide identical keys. This PR adds evidence that verl's own built-in scorer return shapes can differ, applies one shared implementation to both framework paths, and covers schemas split across separate agent-loop workers in both rich-first and poor-first order.
process_validation_metrics() crashes on None-filled sparse reward_extra_info keys (metrics emitted for only some samples) #6830 / [trainer] Fix process_validation_metrics crash on None-filled sparse reward keys #6845 concern the later aggregation of None values in process_validation_metrics. This PR deliberately does not change metric aggregation; it fixes the earlier batch-assembly crash/data-loss boundary and leaves downstream None handling to [trainer] Fix process_validation_metrics crash on None-filled sparse reward keys #6845.

Checklist Before Starting

Searched open and closed PRs/issues for reward_extra_info, reward_extra_keys, compute_rm_score, and sparse-key handling; compared [reward] fix: union reward_extra_info keys across samples in compute_rm_score #6545 and [trainer] Fix process_validation_metrics crash on None-filled sparse reward keys #6845 as described above.
Format the PR title as [{modules}] {type}: {description}.

Test

CPU-only; no GPU or running Ray cluster required:

pytest \
  tests/experimental/reward_loop/test_reward_extra_info_assembly_on_cpu.py \
  tests/experimental/reward_loop/test_reward_extra_info_call_site_on_cpu.py \
  tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py \
  tests/experimental/fully_async_policy/test_reward_extra_info_on_cpu.py \
  -q
# 17 passed

ruff check <the eight changed files>
# All checks passed!

ruff format --check <the eight changed files>
# 8 files already formatted

Coverage includes empty and uniform schemas, sparse keys in either sample order, explicit None, stable key ordering, dense/sparse dtypes, real GSM8K/math-DAPO scorer outputs, reward-loop assembly, cross-worker agent-loop assembly, fully-async finalization, and collision rejection.

API and behavior

No public API or configuration changes. Uniform-schema batches keep the same keys and inferred dtypes. Heterogeneous batches now retain every emitted field and use None for missing values instead of crashing or dropping data.

Design & Code Changes

verl/utils/reward_score/reward_extra_info.py: shared ordered-union assembler.
verl/experimental/reward_loop/reward_loop.py: use the shared assembler in compute_rm_score.
verl/experimental/agent_loop/agent_loop.py: carry raw metadata through worker chunks and assemble it globally after DataProto.concat.
verl/experimental/fully_async_policy/fully_async_rollouter.py: apply the same finalizer when returning a worker result directly.
CPU regression tests cover helper semantics, built-in scorer schemas, and both call sites.

Checklist Before Submitting

Read the contribution guide and repository AGENTS.md.
Ran focused unit tests and Ruff checks after merging current main.
Added CPU regression tests.
Re-request maintainer CI through the contributor Feishu channel after this refresh.
AI assistance was used (Claude and OpenAI Codex). The assistance and verification commands are disclosed above.\n- [ ] I have personally reviewed every changed line and can explain the implementation and its scope.

… loop `RewardLoopManager.compute_rm_score` inferred the extra-info schema from sample 0 only (`list(reward_extra_infos[0].keys())`) and then indexed every sample with `info[key]`. Per-sample schemas legitimately differ: * reward functions emit diagnostic keys conditionally (e.g. only on parse success) — the same pattern the v1 trainer already null-fills for at verl/trainer/ppo/v1/trainer_base.py:1052-1058; * `default_compute_score` returns a dict for some data sources (`math_dapo`) and a bare float for others (`openai/gsm8k`), so any mixed-dataset batch has two schemas ({"score","acc","pred"} vs {"acc"}). Consequences depend on which sample lands at index 0 after chunking across reward workers: * richer sample first -> `KeyError` crash in `compute_rm_score` * poorer sample first -> the extra columns are silently dropped from `non_tensor_batch` / `reward_extra_keys` (missing from dumps and metrics) Assemble the batch from the union of keys instead, filling samples that did not emit a key with `None` (object dtype only when a key is sparse; dense keys keep their natural numpy dtype). Logic is extracted into a module-level `assemble_reward_extra_info` helper with CPU regression tests covering both failure directions. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>

gemini-code-assist · 2026-07-26T10:47:13Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

CLAassistant · 2026-07-26T10:47:16Z

All committers have signed the CLA.

Bring in upstream CI and FSDP test-fixture fixes without rewriting the public PR branch history. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

Use one ordered-union assembler in both reward-loop batch paths and cover heterogeneous built-in scorer schemas and call sites.\n\nCo-authored-by: OpenAI Codex <[email protected]>

Refresh the public PR branch against current upstream without rewriting history.\n\nCo-authored-by: OpenAI Codex <[email protected]>

Carry raw reward metadata through worker boundaries, finalize it in regular and fully-async managers, and reject field collisions before trainer union.\n\nCo-authored-by: OpenAI Codex <[email protected]>

### What does this PR do? Fixes #7486. `AgentLoopOutput.as_dict()` uses `model_dump(exclude_unset=True)`, which omits `extra_fields` when callers rely on the declared default. Mutating that default mapping after construction does not add the field to Pydantic's fields-set, so the existing teacher-field promotion raises `KeyError` and dynamic fields are lost. The fix restores a shallow copy of `self.extra_fields` only when the serialized output omitted the field, then performs the existing teacher-field promotion on that output mapping. Explicitly supplied `extra_fields` continue to use Pydantic's serialized mapping, and serialization does not mutate the model. This is not a duplicate. No open PR is linked to #7486 or matched the `AgentLoopOutput extra_fields exclude_unset` search. #7151 handles sparse reward metadata across samples, and #7019 handles an `output` argument-name collision; neither changes default-field serialization. [Issue search](https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+7486+in%3Abody) · [Keyword search](https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+AgentLoopOutput+extra_fields+exclude_unset) AI assistance disclosure: OpenAI Codex was used to inspect the issue, implement the change, run checks, and draft this PR. The PR remains a draft pending human review of every changed line. ### Checklist Before Starting - [x] Search for similar PRs. Queries and related-PR distinctions are documented above. - [x] Format the PR title as `[{modules}] {type}: {description}`. ### Test - `pre-commit run --all-files --show-diff-on-failure --color=always`: all 13 hooks passed, including Ruff, Ruff format, mypy, generated-config verification, sanity checks, and compileall. - Prior Model Factory CPU agent-loop regression suite containing the mutated-default case: 8 tests passed. - Model Factory Qwen3-4B LangGraph GRPO `val_only` run on 2 H20 GPUs: platform phase 5, container exit code 0, MATH `acc/mean@1=0.996`, and no `KeyError: 'extra_fields'` or Python traceback in the full job log. - The exact three-case test file is automatically collected by the existing `cpu_unit_tests.yml` `*_on_cpu.py` pattern; its PR check is expected to run on this branch. ### API and Usage Example No public API or configuration changes. Valid `AgentLoopOutput` instances that omit `extra_fields` now serialize successfully, and fields added after construction are retained. ### Design & Code Changes - Use `output.setdefault("extra_fields", self.extra_fields.copy())` so only an omitted default is restored. - Preserve the existing promotion of `teacher_ids` and `teacher_logprobs` without mutating `self.extra_fields`. - Add CPU regression coverage for an untouched default, a dynamically populated default, and teacher-field promotion with model-state preservation. ### Checklist Before Submitting - [x] Read the contribution guide and current `AGENTS.md`. - [x] Apply the complete pre-commit suite. - [ ] Add / Update documentation. Not needed because there is no user-facing API or configuration change. - [x] Add CPU regression tests covered by the existing CPU CI workflow. - [ ] Request project CI in Slack or Feishu after human review if the automatic checks require maintainer approval. - [ ] Update the recipe submodule reference. Not applicable because this PR does not modify the recipe submodule. Signed-off-by: caozs2 <[email protected]> Co-authored-by: OpenAI Codex <[email protected]>

…7491) ### What does this PR do? Fixes verl-project#7486. `AgentLoopOutput.as_dict()` uses `model_dump(exclude_unset=True)`, which omits `extra_fields` when callers rely on the declared default. Mutating that default mapping after construction does not add the field to Pydantic's fields-set, so the existing teacher-field promotion raises `KeyError` and dynamic fields are lost. The fix restores a shallow copy of `self.extra_fields` only when the serialized output omitted the field, then performs the existing teacher-field promotion on that output mapping. Explicitly supplied `extra_fields` continue to use Pydantic's serialized mapping, and serialization does not mutate the model. This is not a duplicate. No open PR is linked to verl-project#7486 or matched the `AgentLoopOutput extra_fields exclude_unset` search. verl-project#7151 handles sparse reward metadata across samples, and verl-project#7019 handles an `output` argument-name collision; neither changes default-field serialization. [Issue search](https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+7486+in%3Abody) · [Keyword search](https://github.com/verl-project/verl/pulls?q=is%3Apr+is%3Aopen+AgentLoopOutput+extra_fields+exclude_unset) AI assistance disclosure: OpenAI Codex was used to inspect the issue, implement the change, run checks, and draft this PR. The PR remains a draft pending human review of every changed line. ### Checklist Before Starting - [x] Search for similar PRs. Queries and related-PR distinctions are documented above. - [x] Format the PR title as `[{modules}] {type}: {description}`. ### Test - `pre-commit run --all-files --show-diff-on-failure --color=always`: all 13 hooks passed, including Ruff, Ruff format, mypy, generated-config verification, sanity checks, and compileall. - Prior Model Factory CPU agent-loop regression suite containing the mutated-default case: 8 tests passed. - Model Factory Qwen3-4B LangGraph GRPO `val_only` run on 2 H20 GPUs: platform phase 5, container exit code 0, MATH `acc/mean@1=0.996`, and no `KeyError: 'extra_fields'` or Python traceback in the full job log. - The exact three-case test file is automatically collected by the existing `cpu_unit_tests.yml` `*_on_cpu.py` pattern; its PR check is expected to run on this branch. ### API and Usage Example No public API or configuration changes. Valid `AgentLoopOutput` instances that omit `extra_fields` now serialize successfully, and fields added after construction are retained. ### Design & Code Changes - Use `output.setdefault("extra_fields", self.extra_fields.copy())` so only an omitted default is restored. - Preserve the existing promotion of `teacher_ids` and `teacher_logprobs` without mutating `self.extra_fields`. - Add CPU regression coverage for an untouched default, a dynamically populated default, and teacher-field promotion with model-state preservation. ### Checklist Before Submitting - [x] Read the contribution guide and current `AGENTS.md`. - [x] Apply the complete pre-commit suite. - [ ] Add / Update documentation. Not needed because there is no user-facing API or configuration change. - [x] Add CPU regression tests covered by the existing CPU CI workflow. - [ ] Request project CI in Slack or Feishu after human review if the automatic checks require maintainer approval. - [ ] Update the recipe submodule reference. Not applicable because this PR does not modify the recipe submodule. Signed-off-by: caozs2 <[email protected]> Co-authored-by: OpenAI Codex <[email protected]>

keepkeen requested review from ArronHZG and wuxibin89 as code owners July 26, 2026 10:47

keepkeen and others added 4 commits July 29, 2026 16:37

Merge latest main into fix/reward-loop-sparse-extra-info

1a54755

Bring in upstream CI and FSDP test-fixture fixes without rewriting the public PR branch history. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>

[reward, agent_loop] fix: share sparse extra-info assembly

d0a60fe

Use one ordered-union assembler in both reward-loop batch paths and cover heterogeneous built-in scorer schemas and call sites.\n\nCo-authored-by: OpenAI Codex <[email protected]>

Merge latest main into fix/reward-loop-sparse-extra-info

60ad74b

Refresh the public PR branch against current upstream without rewriting history.\n\nCo-authored-by: OpenAI Codex <[email protected]>

[agent_loop] fix: finalize reward metadata after dispatch

96eead6

Carry raw reward metadata through worker boundaries, finalize it in regular and fully-async managers, and reject field collisions before trainer union.\n\nCo-authored-by: OpenAI Codex <[email protected]>

keepkeen changed the title ~~[reward] fix: sparse reward_extra_info keys crash or vanish in reward loop batch assembly~~ [reward, rollout] fix: align sparse reward_extra_info across batch paths Aug 5, 2026

adityasingh2400 mentioned this pull request Aug 5, 2026

[reward] fix: do not drop a GSM8k answer that is followed by more text #7274

Open

7 of 8 tasks

ai-yang mentioned this pull request Aug 12, 2026

Standalone DisRM crashes async agent reward scoring when reward_extra_info is missing #7368

Open

3 tasks done

le-czs mentioned this pull request Aug 20, 2026

[rollout] fix: preserve default AgentLoop extra fields #7491

Merged

5 of 8 tasks

jmaicher mentioned this pull request Aug 27, 2026

[trainer] fix: dump reward_extra_info in v1 train rollout JSONL #7588

Open

5 of 8 tasks

amanyagami mentioned this pull request Aug 29, 2026

[reward] fix: DisRM returns reward_extra_info so async scoring does not KeyError #7608

Open

6 tasks done

feiiiiii5 mentioned this pull request Nov 3, 2026

Preprint release: verifier-failure audit, repaired chain, offline policy analysis, current paper and arXiv source feiiiiii5/rewardfail#1

Merged

Navigation Menu

Repository navigation

[reward, rollout] fix: align sparse reward_extra_info across batch paths - #7151

[reward, rollout] fix: align sparse reward_extra_info across batch paths#7151
keepkeen wants to merge 5 commits into
verl-project:mainfrom
keepkeen:fix/reward-loop-sparse-extra-info

keepkeen commented Jul 26, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Jul 26, 2026

Uh oh!

CLAassistant commented Jul 26, 2026 •

edited

Loading

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

keepkeen commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Relationship to earlier work

Checklist Before Starting

Test

API and behavior

Design & Code Changes

Checklist Before Submitting

Uh oh!

gemini-code-assist Bot commented Jul 26, 2026

Uh oh!

CLAassistant commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

keepkeen commented Jul 26, 2026 •

edited

Loading

CLAassistant commented Jul 26, 2026 •

edited

Loading