初始化项目，由ModelHub XC社区提供模型

Model: bk1dr/qwen3-8b-code-pkpo
Source: Original Platform

pkpo.py

@@ -0,0 +1,98 @@
+"""PKPO reward transformations — Listing 1 of the paper (arXiv:2505.15201), verbatim.
+
+sloo_minus_one is the s^(loo-1) estimator of Eq. (33)/(34): transformed rewards whose
+sum estimates pass@k / max_g@k, with a k-1-subset LOO baseline. Requires n >= k+1.
+At k=1 the paper uses untransformed rewards with mean centering (its k-1 baseline is
+undefined); transform_rewards() below handles that case.
+"""
+from typing import Callable
+
+import numpy as np
+
+
+def _m_normed(N: int, K: int, i: int, j: int) -> float:
+    if i == j and i >= K - 1:
+        return (
+            K / (N - K + 1)
+            * np.prod(np.arange(i - K + 2, i + 1) / np.arange(N - K + 2, N + 1))
+        )
+    elif j > i and j >= K - 1 and K >= 2:
+        return (
+            K / (N - K + 1) * (K - 1) / N
+            * np.prod(np.arange(j - K + 2, j) / np.arange(N - K + 2, N))
+        )
+    return 0
+
+
+def _m_diagonal(N: int, K: int) -> np.ndarray:
+    return np.array([_m_normed(N, K, i, i) for i in range(N)])
+
+
+def rho(g: np.ndarray, K: int) -> float:
+    """See Equation (12)."""
+    return (np.sort(g) * _m_diagonal(len(g), K)).sum()
+
+
+def _delta(N: int, K: int, i: int) -> float:
+    return _m_normed(N, K, i, i + 1) - _m_normed(N, K, i + 1, i + 1)
+
+
+def _deltas(N: int, K: int) -> np.ndarray:
+    return np.array([_delta(N - 1, K, i) for i in range(N - 2)])
+
+
+def _sorted_apply(func: Callable) -> Callable:
+    def inner(x: np.ndarray, *args, **kwargs) -> np.ndarray:
+        i_sort = np.argsort(x)
+        func_x = np.zeros_like(x)
+        func_x[i_sort] = func(x[i_sort], *args, **kwargs)
+        return func_x
+    return inner
+
+
+@_sorted_apply
+def s(g: np.ndarray, K: int):
+    """See Equation (19)."""
+    N = len(g)
+    c = g * _m_diagonal(N, K)
+    c[:(N - 1)] += g[1:] * _deltas(N + 1, K)
+    return np.cumsum(c[::-1])[::-1]
+
+
+@_sorted_apply
+def _b(g: np.ndarray, K: int) -> np.ndarray:
+    N = len(g)
+    w = (_m_diagonal(N - 1, K) * np.arange(1, N)).astype(float)
+    w[1:] += _deltas(N, K) * np.arange(1, N - 1)
+    c1 = np.array([(w * g[1:]).sum()])
+    c2 = (g[:-1] - g[1:]) * w
+    return np.cumsum(np.concatenate((c1, c2)))
+
+
+def sloo(g: np.ndarray, K: int) -> np.ndarray:
+    """See Equation (29)."""
+    return s(g, K) - _b(g, K) / (len(g) - 1)
+
+
+def sloo_minus_one(g: np.ndarray, K: int) -> np.ndarray:
+    """See Equation (33)."""
+    return s(g, K) - _b(g, K - 1) * K / (K - 1) / len(g)
+
+
+def transform_rewards(g: np.ndarray, K: int) -> np.ndarray:
+    """PKPO advantages for one group of n rollouts of the same prompt.
+
+    k >= 2: sloo_minus_one exactly as in Listing 1.
+    k == 1: s(g, 1) (= g/n, no transformation) with group-mean centering, which is
+    what the paper uses for its k_opt=1 runs ("without which the training diverges").
+    Scaled by n so advantage magnitude is O(reward) at every k (constant across
+    stages; equivalent to a learning-rate rescale).
+    """
+    g = np.asarray(g, dtype=np.float64)
+    n = len(g)
+    if K >= 2:
+        out = sloo_minus_one(g, K)
+    else:
+        out = s(g, 1)
+        out = out - out.mean()
+    return out * n