check_reward_nonzero_std drops every group when n_samples_per_prompt=1, hanging rollout

Open Beginner friendly
#252 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
1/5
Estimated time
1-3 hours
Newbie friendliness
88/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python, pytorch

Research direction

Start in miles/rollout/filter_hub/dynamic_sampling_filters.py at check_reward_nonzero_std, then inspect how generate_rollout_async in sglang_diffusion_rollout.py handles rejected groups. Reproduce the single-reward torch.std behavior and verify that a one-sample group is retained, the returned keep value is a plain bool, and rollout no longer hangs when n_samples_per_prompt is 1.

Written by the indexing model from the issue text.

Description

Bug

check_reward_nonzero_std in miles/rollout/filter_hub/dynamic_sampling_filters.py always drops groups made of a single sample.

def check_reward_nonzero_std(args, samples: list[Sample], **kwargs):
    rewards = [sample.get_reward_value(args) for sample in samples]
    keep = torch.tensor(rewards, dtype=torch.float).std() > 0.0
    ...

torch.std() on a 1-element tensor divides by (n - 1) == 0 and returns NaN. NaN > 0.0 evaluates to False in PyTorch, so keep is always False regardless of the actual reward.

Repro:

import torch
rewards = [5.0]
torch.tensor(rewards, dtype=torch.float).std() > 0.0
# tensor(False)  (with a UserWarning about degrees of freedom <= 0)

Impact

With --n-samples-per-prompt 1 (a supported configuration) combined with this filter (recommended as the example dynamic-sampling filter in arguments.py), every rollout group is rejected. The rollout loop (generate_rollout_async in sglang_diffusion_rollout.py) keeps resubmitting generation requests forever since nothing is ever accepted — a silent infinite hang that burns rollout compute instead of failing clearly.

Fix

Special-case len(rewards) <= 1 to always keep the sample (there's no variance to check with one sample), and cast keep to a plain bool instead of leaving it as a torch.Tensor.

PR: (linked)

Dominant language
Python
Stars
96
Forks
17
Avg merge
1d 10h
Merged PRs (30d)
31

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from radixark/miles_diffusion

All issues in radixark/miles_diffusion

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.