1import json
 2import numpy as np
 3from diversity_baseline import diversity_baseline, uniform_loo, effective_count
 4
 5
 6def main():
 7    out = {}
 8    # Prediction 1: p=0 is exactly uniform LOO, independent of embeddings.
 9    rng = np.random.default_rng(2183)
10    costs = rng.normal(size=6)
11    emb = rng.normal(size=(6, 3))
12    b0, a0, q0, _ = diversity_baseline(costs, emb, p=0)
13    bu = uniform_loo(costs)
14    out['p0_max_baseline_error'] = float(np.max(np.abs(b0 - bu)))
15    out['p0_effective_count_observed_predicted'] = [float(effective_count(q0).mean()), 5.0]
16
17    # Prediction 2: for two peer distances, weight ratio is (d_far/d_near)^p.
18    # eps is negligible here; use a group with one anchor and two peers.
19    costs2 = np.array([0., 1., 2.])
20    emb2 = np.array([[0., 0.], [1., 0.], [4., 0.]])
21    _, _, q2, d2 = diversity_baseline(costs2, emb2, p=1, eps=1e-8)
22    observed_ratio = q2[0, 2] / q2[0, 1]
23    predicted_ratio = (d2[0, 2] + 1e-8) / (d2[0, 1] + 1e-8)
24    out['distance_ratio_observed_predicted'] = [float(observed_ratio), float(predicted_ratio)]
25
26    # Prediction 3: duplicated/near peers are downweighted; sweep p and compare
27    # effective count with the exact two-distance prediction for the anchor.
28    # Anchor at 0, two near peers at +/-delta and two far peers at +/-D.
29    delta, D = 0.1, 2.0
30    emb3 = np.array([[0., 0.], [delta, 0.], [-delta, 0.], [D, 0.], [-D, 0.]])
31    costs3 = np.arange(5, dtype=float)
32    sweep = []
33    for p in [0., .5, 1., 2.]:
34        _, _, q, d = diversity_baseline(costs3, emb3, p=p)
35        near = (d[0, 1] + 1e-8) ** p
36        far = (d[0, 3] + 1e-8) ** p
37        predicted_eff = 1.0 / (2 * (near/(2*near+2*far))**2 + 2 * (far/(2*near+2*far))**2)
38        sweep.append({'p': p, 'effective_observed': float(effective_count(q)[0]),
39                      'effective_predicted': float(predicted_eff),
40                      'far_to_near_weight_observed': float(q[0,3]/q[0,1]),
41                      'far_to_near_weight_predicted': float(far/near)})
42    out['effective_count_sweep'] = sweep
43
44    # Small grouped-policy proxy: categorical logits, score-function gradients.
45    # Actions 0/1 are near-duplicates and cheap; 2/3 are structurally distinct
46    # and expensive. Diversity weighting should alter, not automatically improve,
47    # gradient variance; this is an honest secondary comparison.
48    probs = np.array([.35, .35, .15, .15])
49    action_emb = np.array([[0.,0.], [.02,0.], [1.,0.], [1.02,0.]])
50    action_cost = np.array([1., 1.1, 4., 4.1])
51    n_groups, B = 20000, 8
52    grads_u, grads_d = [], []
53    for _ in range(n_groups):
54        acts = rng.choice(4, size=B, p=probs)
55        c = action_cost[acts]
56        z = action_emb[acts]
57        bu = uniform_loo(c)
58        bd, ad, _, _ = diversity_baseline(c, z, p=1.)
59        au = bu - c
60        score = np.eye(4)[acts] - probs
61        grads_u.append(((-au[:,None] * score)).mean(0))
62        grads_d.append(((-ad[:,None] * score)).mean(0))
63    var_u = np.var(grads_u, axis=0, ddof=1).mean()
64    var_d = np.var(grads_d, axis=0, ddof=1).mean()
65    out['gradient_variance_uniform_diversity_ratio'] = [float(var_u), float(var_d), float(var_d/var_u)]
66    print(json.dumps(out, indent=2))
67
68if __name__ == '__main__':
69    main()