File size: 5,059 Bytes
53c10a4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """Pure helpers for OraRL post-selection moment correction."""
import math
from dataclasses import dataclass
import torch
SIGMA_OP_FALLBACK_THRESHOLD = 1e-3
@dataclass(frozen=True)
class PostSelectionReference:
"""Pre-selection on-policy statistics for one rollout group."""
on_policy_rms: float
sigma_op: float
on_policy_rows: int = 0
def balance_post_selection_group(
active_advantages: torch.Tensor,
is_oracle_row: torch.Tensor,
*,
reference: PostSelectionReference,
recenter: bool,
rms_match: bool,
rms_min_scale: float = 0.25,
eps: float = 1e-6,
) -> tuple[torch.Tensor, dict[str, float]]:
"""Recenter one selected group and optionally match its pre-selection RMS.
The active group contains the policy rows retained by OraRL selection and
one detached oracle row. If recentering would make the oracle advantage
negative, the vector is projected onto ``sum(A)=0, A_oracle>=0`` by setting
the oracle to zero and distributing the correction evenly over policy rows.
RMS matching can only downscale.
"""
if active_advantages.ndim != 1:
raise ValueError(
"active_advantages must be one-dimensional, got "
f"{tuple(active_advantages.shape)}."
)
if (
is_oracle_row.ndim != 1
or is_oracle_row.shape != active_advantages.shape
):
raise ValueError(
"is_oracle_row must match active_advantages, got "
f"{tuple(is_oracle_row.shape)} and {tuple(active_advantages.shape)}."
)
if active_advantages.numel() == 0:
raise ValueError("active_advantages must not be empty.")
if not bool(torch.isfinite(active_advantages).all()):
raise ValueError("active_advantages contains non-finite values.")
if rms_match and not recenter:
raise ValueError("rms_match requires recenter.")
if not 0.0 <= rms_min_scale <= 1.0:
raise ValueError(
f"rms_min_scale must be in [0, 1], got {rms_min_scale}."
)
if (
not math.isfinite(reference.on_policy_rms)
or not math.isfinite(reference.sigma_op)
or reference.on_policy_rms < 0.0
or reference.sigma_op < 0.0
):
raise ValueError(
"reference RMS and sigma_op must be finite and non-negative."
)
oracle_mask = is_oracle_row.to(
device=active_advantages.device,
dtype=torch.bool,
)
oracle_rows = int(oracle_mask.sum().item())
if oracle_rows > 1:
raise ValueError(
"post-selection correction supports at most one oracle row, "
f"got {oracle_rows}."
)
policy_rows = int((~oracle_mask).sum().item())
if oracle_rows == 1 and policy_rows == 0:
raise ValueError(
"post-selection correction requires at least one active policy row "
"alongside the oracle row."
)
active_before = active_advantages
mean_before = active_before.mean()
rms_before = torch.sqrt(torch.mean(active_before.square()))
active_after = active_before.clone()
oracle_projection = active_before.new_zeros(())
if recenter:
active_after -= mean_before
if (
oracle_rows == 1
and float(active_after[oracle_mask].item()) < 0.0
):
policy_mask = ~oracle_mask
correction = -active_after[oracle_mask]
active_after[oracle_mask] = 0.0
active_after[policy_mask] -= correction / float(policy_rows)
oracle_projection = active_before.new_ones(())
rms_scale = active_before.new_ones(())
sigma_fallback = reference.sigma_op < SIGMA_OP_FALLBACK_THRESHOLD
if rms_match:
active_rms = torch.sqrt(torch.mean(active_after.square()))
if float(active_rms.item()) > eps and not sigma_fallback:
target_rms = active_after.new_tensor(reference.on_policy_rms)
rms_scale = torch.clamp(
target_rms / (active_rms + eps),
min=float(rms_min_scale),
max=1.0,
)
active_after *= rms_scale
mean_after = active_after.mean()
rms_after = torch.sqrt(torch.mean(active_after.square()))
oracle_after = (
float(active_after[oracle_mask].item())
if oracle_rows == 1
else 0.0
)
metrics = {
"active_mean_before": float(mean_before.item()),
"active_rms_before": float(rms_before.item()),
"on_policy_rms": float(reference.on_policy_rms),
"rms_scale": float(rms_scale.item()),
"oracle_sign_projection": float(oracle_projection.item()),
"sigma_op_fallback": float(sigma_fallback),
"active_mean_after": float(mean_after.item()),
"active_rms_after": float(rms_after.item()),
"oracle_advantage_after": oracle_after,
"active_rows": float(active_after.numel()),
"active_policy_rows": float((~oracle_mask).sum().item()),
"oracle_rows": float(oracle_rows),
}
return active_after, metrics
|