code/verl/trainer/orarl_selection.py

"""OraRL's strict sign-balanced rollout selection."""

from collections import defaultdict
from typing import Any

import numpy as np
import torch

from ..protocol import DataProto
from ..utils import torch_functional as VF


def _sequence_advantages(
    advantages: torch.Tensor,
    response_mask: torch.Tensor,
) -> torch.Tensor:
    if advantages.shape != response_mask.shape:
        raise ValueError(
            "advantages and response_mask must have identical shapes, got "
            f"{tuple(advantages.shape)} and {tuple(response_mask.shape)}."
        )
    return VF.masked_mean(advantages, response_mask, dim=-1)


def _rank_by_magnitude(rows: list[int], scores: torch.Tensor) -> list[int]:
    return sorted(
        rows,
        key=lambda row: (-abs(float(scores[row].item())), row),
    )


@torch.no_grad()
def select_orarl_rollouts(
    data: DataProto,
    *,
    n_rollouts: int,
    prune_ratio: float,
    world_size: int,
    positive_quota: int,
    negative_quota: int,
    oracle_key: str = "is_oracle_row",
) -> tuple[DataProto, dict[str, float]]:
    """Retain one oracle row and sign-balanced policy rows in every group."""

    if not 0.0 < prune_ratio < 1.0:
        raise ValueError("OraRL prune_ratio must be in (0, 1).")
    if n_rollouts <= 1:
        raise ValueError("OraRL n_rollouts must be greater than one.")
    if world_size <= 0:
        raise ValueError("OraRL world_size must be positive.")
    keep_per_group = max(1, int(n_rollouts * (1.0 - prune_ratio)))
    if positive_quota < 0 or negative_quota < 0:
        raise ValueError("OraRL sign quotas must be non-negative.")
    if positive_quota + negative_quota + 1 != keep_per_group:
        raise ValueError(
            "OraRL sign quotas plus one oracle must equal the keep budget."
        )

    required_batch = ("advantages", "response_mask", "attention_mask")
    missing_batch = sorted(
        key for key in required_batch if key not in data.batch
    )
    if missing_batch:
        raise ValueError(
            "OraRL selection is missing batch fields: "
            + ", ".join(missing_batch)
        )
    if "uid" not in data.non_tensor_batch:
        raise ValueError("OraRL selection requires non_tensor_batch['uid'].")
    if oracle_key not in data.non_tensor_batch:
        raise ValueError(
            f"OraRL selection requires non_tensor_batch[{oracle_key!r}]."
        )

    signed_scores = _sequence_advantages(
        data.batch["advantages"],
        data.batch["response_mask"],
    )
    total_rows = int(signed_scores.numel())
    uids = np.asarray(data.non_tensor_batch["uid"], dtype=object)
    oracle_flags = np.asarray(
        data.non_tensor_batch[oracle_key],
        dtype=bool,
    )
    if len(uids) != total_rows or len(oracle_flags) != total_rows:
        raise ValueError(
            "OraRL uid and oracle flags must align with batch rows."
        )

    grouped: dict[Any, list[int]] = defaultdict(list)
    for row, uid in enumerate(uids):
        grouped[uid].append(row)

    selected_rows: list[int] = []
    positive_kept = 0
    negative_kept = 0
    zero_kept = 0
    cross_sign_fallback = 0
    zero_fallback = 0

    for uid, rows in grouped.items():
        oracle_rows = [row for row in rows if bool(oracle_flags[row])]
        if len(oracle_rows) != 1:
            raise ValueError(
                f"OraRL group {uid!r} requires exactly one oracle row, "
                f"got {len(oracle_rows)}."
            )
        candidates = [row for row in rows if row != oracle_rows[0]]
        if len(candidates) != n_rollouts:
            raise ValueError(
                f"OraRL group {uid!r} requires {n_rollouts} policy rows, "
                f"got {len(candidates)}."
            )

        positive = _rank_by_magnitude(
            [
                row
                for row in candidates
                if float(signed_scores[row].item()) > 0.0
            ],
            signed_scores,
        )
        negative = _rank_by_magnitude(
            [
                row
                for row in candidates
                if float(signed_scores[row].item()) < 0.0
            ],
            signed_scores,
        )
        zeros = [
            row
            for row in candidates
            if float(signed_scores[row].item()) == 0.0
        ]

        chosen = positive[:positive_quota] + negative[:negative_quota]
        chosen_set = set(chosen)
        remaining = keep_per_group - 1 - len(chosen)
        if remaining > 0:
            surplus = _rank_by_magnitude(
                [
                    row
                    for row in positive + negative
                    if row not in chosen_set
                ],
                signed_scores,
            )
            cross_fill = surplus[:remaining]
            chosen.extend(cross_fill)
            chosen_set.update(cross_fill)
            cross_sign_fallback += len(cross_fill)
            remaining -= len(cross_fill)
        if remaining > 0:
            zero_fill = [
                row for row in zeros if row not in chosen_set
            ][:remaining]
            chosen.extend(zero_fill)
            chosen_set.update(zero_fill)
            zero_fallback += len(zero_fill)
            remaining -= len(zero_fill)
        if remaining:
            raise RuntimeError(
                f"OraRL group {uid!r} could not satisfy its keep budget."
            )

        for row in chosen:
            score = float(signed_scores[row].item())
            positive_kept += int(score > 0.0)
            negative_kept += int(score < 0.0)
            zero_kept += int(score == 0.0)
        selected_rows.extend([oracle_rows[0], *chosen])

    selected_rows.sort()
    if len(selected_rows) % world_size:
        raise RuntimeError(
            f"OraRL selected batch size {len(selected_rows)} is not "
            f"divisible by world_size {world_size}."
        )

    selected = data[selected_rows]
    selected.meta_info = dict(selected.meta_info)
    selected.meta_info["global_token_num"] = (
        torch.sum(selected.batch["attention_mask"], dim=-1).tolist()
    )

    selected_index = torch.tensor(
        selected_rows,
        dtype=torch.long,
        device=signed_scores.device,
    )
    selected_set = set(selected_rows)
    dropped_rows = [
        row for row in range(total_rows) if row not in selected_set
    ]
    metrics = {
        "orarl/selection/groups": float(len(grouped)),
        "orarl/selection/keep_per_group": float(keep_per_group),
        "orarl/selection/kept_rows": float(len(selected_rows)),
        "orarl/selection/dropped_rows": float(
            total_rows - len(selected_rows)
        ),
        "orarl/selection/effective_keep_ratio": (
            len(selected_rows) / float(total_rows)
        ),
        "orarl/selection/oracle_rows_forced": float(len(grouped)),
        "orarl/selection/positive_policy_rows_kept": float(positive_kept),
        "orarl/selection/negative_policy_rows_kept": float(negative_kept),
        "orarl/selection/zero_policy_rows_kept": float(zero_kept),
        "orarl/selection/cross_sign_fallback_rows": float(
            cross_sign_fallback
        ),
        "orarl/selection/zero_fallback_rows": float(zero_fallback),
        "orarl/selection/abs_advantage_kept_mean": float(
            signed_scores.index_select(0, selected_index).abs().mean().item()
        ),
    }
    if dropped_rows:
        dropped_index = torch.tensor(
            dropped_rows,
            dtype=torch.long,
            device=signed_scores.device,
        )
        metrics["orarl/selection/abs_advantage_dropped_mean"] = float(
            signed_scores.index_select(0, dropped_index).abs().mean().item()
        )
    return selected, metrics