Skip to content

fix(rollout): reject aborted samples before reward tensor - #2361

Open
YeonwooSung wants to merge 1 commit into
THUDM:mainfrom
YeonwooSung:fix/200-reject-aborted-train-rewards
Open

fix(rollout): reject aborted samples before reward tensor#2361
YeonwooSung wants to merge 1 commit into
THUDM:mainfrom
YeonwooSung:fix/200-reject-aborted-train-rewards

Conversation

@YeonwooSung

Copy link
Copy Markdown

Fixes #200.

Aborted samples keep reward=None. On the default GRPO path that list was passed to torch.tensor, which crashes. Custom generate functions that abort (e.g. Search-R1 tool failure) hit this.

This PR fail-closes instead of scoring those samples. We do not default reward to 0.0. A previous attempt that did (#1252) was closed unmerged; aborted / unrewarded samples must not reach training.

Change

In RolloutManager._post_process_rewards, on the default path only (when custom_reward_post_process_func is None), reject any sample with status == Sample.Status.ABORTED or reward is None before get_reward_value / torch.tensor.

The error tells callers to filter in the generate function or recycle via partial rollout. Custom post-process keeps its own contract and is not subject to this check.

get_reward_value is self.reward or self.reward[args.reward_key]. Checking reward is None first avoids a TypeError when reward_key is set.

Tests

tests/test_post_process_rewards.py (NUM_GPUS = 0) binds _post_process_rewards onto a dummy object (no Ray):

  1. Two completed numeric rewards succeed and return tensor-compatible lists
  2. One ABORTED sample raises ValueError
  3. COMPLETED with reward is None raises ValueError
  4. A custom post-process function is used as-is (aborted sample does not hit the new check)

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.

UX/Unintended bug: sample.reward is None on aborted sample

1 participant