奖励减去同批平均值，为什么学习信号会少一截？

"""Exact policy-gradient expectations for a one-step Bernoulli policy."""
from fractions import Fraction as F
from itertools import product


def estimate(actions, p, method):
    n = len(actions)
    if not n or (method == "leave_one_out" and n < 2):
        raise ValueError("leave-one-out needs at least two samples")
    rewards = list(actions)  # Fixed reward: R(a) = a.
    mean = F(sum(rewards), n)
    total = F(0)
    for action, reward in zip(actions, rewards):
        if method == "none":
            baseline = F(0)
        elif method == "constant":
            baseline = F(1, 2)
        elif method == "batch_mean":
            baseline = mean
        elif method == "leave_one_out":
            baseline = F(sum(rewards) - reward, n - 1)
        else:
            raise ValueError("unknown baseline")
        score = action - p  # d log pi(a) / d logit(theta).
        total += (reward - baseline) * score
    return total / n


def moments(n, p, method):
    mean = second = F(0)
    for actions in product((0, 1), repeat=n):
        count = sum(actions)
        probability = p ** count * (1 - p) ** (n - count)
        value = estimate(actions, p, method)
        mean += probability * value
        second += probability * value * value
    return mean, second - mean * mean


def main():
    p = F(1, 2)
    exact = p * (1 - p)
    checks = {}
    for probability in (F(1, 4), F(1, 2), F(3, 4)):
        truth = probability * (1 - probability)
        for n in (1, 2, 4):
            checks[f"p={probability},N={n},plain"] = moments(n, probability, "none")[0] == truth
            checks[f"p={probability},N={n},centered"] = moments(n, probability, "batch_mean")[0] == F(n - 1, n) * truth
            if n >= 2:
                checks[f"p={probability},N={n},loo"] = moments(n, probability, "leave_one_out")[0] == truth
                checks[f"p={probability},N={n},sample_identity"] = all(
                    estimate(a, probability, "leave_one_out") == F(n, n - 1) * estimate(a, probability, "batch_mean")
                    for a in product((0, 1), repeat=n))
    try:
        estimate((1,), p, "leave_one_out")
    except ValueError:
        checks["N1_rejected"] = True
    else:
        checks["N1_rejected"] = False
    copied = ((1 - p) * estimate((0, 0), p, "leave_one_out")
              + p * estimate((1, 1), p, "leave_one_out"))
    checks["correlated_copy_breaks_unbiasedness"] = copied == 0 and copied != exact
    checks["constant_baseline_zero_variance_in_this_example"] = moments(2, p, "constant")[1] == 0
    checks["loo_can_increase_variance"] = moments(2, p, "leave_one_out")[1] > moments(2, p, "none")[1]
    assert all(checks.values()), checks
    print("policy: P(action=1)=0.5; reward=action; derivative is wrt logit")
    print(f"exact_gradient={float(exact):.6f}")
    for n in (1, 2, 4):
        plain = moments(n, p, "none")[0]
        centered = moments(n, p, "batch_mean")[0]
        loo = "undefined" if n == 1 else f"{float(moments(n, p, 'leave_one_out')[0]):.6f}"
        print(f"N={n} batches={2**n} plain={float(plain):.6f} batch_mean={float(centered):.6f} leave_one_out={loo}")
    print(f"N=2 variance: plain={float(moments(2, p, 'none')[1]):.6f} leave_one_out={float(moments(2, p, 'leave_one_out')[1]):.6f} constant_0.5={float(moments(2, p, 'constant')[1]):.6f}")
    print(f"copied_samples_N=2: leave_one_out={float(copied):.6f}")
    print(f"PASS {len(checks)} exact checks; independent one-step model, no RLHF training")


if __name__ == "__main__":
    main()