feat(rm): allow per-sample custom_rm_path override

[EazyReal]

Summary

Allow each sample metadata record to override the configured custom reward function path, including batched reward dispatch.

Why

Mixed-task batches need to route samples to different reward functions without splitting rollout jobs. Override samples use the single-sample reward contract; unoverridden samples remain batched through the configured global function. Samples without an override keep the existing behavior. This ports the Slime configuration behavior from THUDM/slime#2081.

Testing

pytest -q tests/fast/rollout/rm_hub/test_rm_hub.py

[gemini-code-assist]

Code Review

This pull request introduces support for per-dataset and per-sample custom reward model (RM) paths, allowing individual samples to override the global --custom-rm-path configuration. The changes update the configuration schemas, sample types, metadata merging utilities, and add corresponding unit tests. Feedback on the changes highlights that while async_rm correctly prioritizes the per-sample custom RM path, batched_async_rm does not, which would cause it to route all samples to the global custom RM and ignore per-sample overrides. A grouped batching implementation is suggested to resolve this issue.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

While async_rm now correctly prioritizes sample.custom_rm_path over args.custom_rm_path, batched_async_rm (defined later in this file) does not.

Specifically, in batched_async_rm:

    if args.custom_rm_path is not None:
        # Ensure the custom reward function is implemented in batch mode
        rm_function = load_function(args.custom_rm_path)
        return await rm_function(args, samples, **kwargs)

If a global args.custom_rm_path is configured, batched_async_rm will unconditionally route all samples to that global custom RM, completely ignoring any per-sample sample.custom_rm_path overrides. This breaks the core objective of supporting mixed-task batches with different reward functions.

To fix this, batched_async_rm should group the samples by their effective custom RM path (either sample.custom_rm_path or args.custom_rm_path) and batch them accordingly, falling back to individual async_rm calls for samples that do not use any custom RM.

Here is a suggested implementation for batched_async_rm:

async def batched_async_rm(
    args,
    samples: list[Sample],
    inplace_set_reward_field: bool = False,
    **kwargs,
) -> list[int | float] | None:
    if inplace_set_reward_field:
        rewards = await batched_async_rm(args, samples, **kwargs)
        for sample, reward in zip(samples, rewards, strict=True):
            assert (
                sample.reward is None
            ), f"Overriding sample.reward from {sample.reward} to {reward}, is this intended?"
            sample.reward = reward
        return None

    # Group samples by their effective custom RM path to preserve batching where possible
    indexed_samples = list(enumerate(samples))
    groups = {}
    for idx, sample in indexed_samples:
        path = sample.custom_rm_path or args.custom_rm_path
        groups.setdefault(path, []).append((idx, sample))

    rewards = [None] * len(samples)
    tasks = []

    for path, group in groups.items():
        sub_indices = [item[0] for item in group]
        sub_samples = [item[1] for item in group]

        if path is not None:
            async def run_custom_batch(p=path, indices=sub_indices, s_samples=sub_samples):
                rm_function = load_function(p)
                batch_rewards = await rm_function(args, s_samples, **kwargs)
                for i, r in zip(indices, batch_rewards):
                    rewards[i] = r
            tasks.append(run_custom_batch())
        else:
            async def run_individual_fallback(indices=sub_indices, s_samples=sub_samples):
                individual_tasks = [async_rm(args, s, **kwargs) for s in s_samples]
                batch_rewards = await asyncio.gather(*individual_tasks)
                for i, r in zip(indices, batch_rewards):
                    rewards[i] = r
            tasks.append(run_individual_fallback())

    await asyncio.gather(*tasks)
    return rewards

[EazyReal]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for per-dataset and per-sample custom reward functions (custom_rm_path), allowing them to override the global --custom-rm-path. The changes span across configuration, sample metadata merging, and the reward model hub (rm_hub) to handle these overrides in both single and batched asynchronous reward model evaluations. Unit tests are also added to verify that per-sample custom reward functions take priority. Feedback on the changes suggests refactoring batched_async_rm to avoid nested asyncio.gather calls and to handle potential synchronous custom reward functions gracefully.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The current implementation of batched_async_rm has a few potential issues:

1. Nested asyncio.gather: It nests asyncio.gather calls (e.g., gathering override_rewards and default_rewards where both are already futures returned by asyncio.gather). This adds unnecessary overhead and complexity.
2. Synchronous RM Compatibility: It assumes that the custom batch reward function (rm_function) is always asynchronous. If a user provides a synchronous custom batch reward function, calling it will return a list of rewards directly, and passing this list to asyncio.gather will raise a TypeError.

We can simplify the logic by flattening the asyncio.gather calls and defensively checking if the batch task is awaitable using inspect.isawaitable. This ensures compatibility with both synchronous and asynchronous custom reward functions while maintaining optimal concurrency.

    samples_without_overrides = [(index, sample) for index, sample in enumerate(samples) if not sample.custom_rm_path]
    override_tasks = [async_rm(args, sample, **kwargs) for _, sample in samples_with_overrides]
    if args.custom_rm_path is not None and samples_without_overrides:
        rm_function = load_function(args.custom_rm_path)
        batch_task = rm_function(args, [sample for _, sample in samples_without_overrides], **kwargs)
        import inspect
        if inspect.isawaitable(batch_task):
            results = await asyncio.gather(*override_tasks, batch_task)
            override_rewards = results[:len(override_tasks)]
            default_rewards = results[len(override_tasks)]
        else:
            results = await asyncio.gather(*override_tasks)
            override_rewards = results
            default_rewards = batch_task
    else:
        non_override_tasks = [async_rm(args, sample, **kwargs) for _, sample in samples_without_overrides]
        results = await asyncio.gather(*override_tasks, *non_override_tasks)
        override_rewards = results[:len(override_tasks)]
        default_rewards = results[len(override_tasks):]

    rewards_by_index = dict(
        zip(
            [index for index, _ in samples_with_overrides] + [index for index, _ in samples_without_overrides],
            [*override_rewards, *default_rewards],
            strict=True,
        )
    )
    return [rewards_by_index[index] for index in range(len(samples))]