1import json, random
  2import numpy as np
  3import torch
  4import torch.nn.functional as F
  5
  6SEED = 1215
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9
 10
 11def js(p, q):
 12    m = (p + q) / 2
 13    return 0.5 * ((p * (p.log() - m.log())).sum(-1) +
 14                  (q * (q.log() - m.log())).sum(-1))
 15
 16
 17def exact_math_check():
 18    cases = [({0, 1}, {0}, True), ({0}, {1}, False),
 19             ({0, 2}, {2}, True), ({1}, {2}, False)]
 20    intersection_ok = all(bool(a & b) == expected for a, b, expected in cases)
 21    p = torch.tensor([[.9, .1], [.9, .1], [.2, .8]], dtype=torch.float64)
 22    return {
 23        "intersection_equivalence": intersection_ok,
 24        "js_zero_equal": float(js(p[0], p[1])),
 25        "js_positive_unequal": float(js(p[0], p[2])),
 26        "js_symmetric_error": abs(float(js(p[0], p[2]) - js(p[2], p[0]))),
 27        "js_upper_bound_log2": float(np.log(2) - js(p[0], p[2])),
 28    }
 29
 30
 31def safety_structure_check():
 32    feasible = [{0}, {0, 1}]
 33    empty = [{0}, {1}]
 34    return {
 35        "feasible_intersection": sorted(set.intersection(*map(set, feasible))),
 36        "empty_intersection": sorted(set.intersection(*map(set, empty))),
 37        "empty_intersection_rate": float(len(set.intersection(*map(set, empty))) == 0),
 38    }
 39
 40
 41def train(consistency_weight, safety_weight, steps=1200):
 42    # Histories 0,1 share observation o_feasible and have compatible safety.
 43    # Histories 2,3 share o_empty and have contradictory safety requirements.
 44    # Imitation labels intentionally conflict within both observation classes.
 45    logits = torch.nn.Parameter(torch.zeros(4, 2))
 46    opt = torch.optim.Adam([logits], lr=0.035)
 47    labels = torch.tensor([0, 1, 0, 1])
 48    margin = 0.90
 49    for _ in range(steps):
 50        q = logits.softmax(-1)
 51        task = F.cross_entropy(logits, labels)
 52        consistency = (js(q[0], q[1]) + js(q[2], q[3])) / 2
 53        # C(x,U) is 1 iff U is safe for x. For the feasible group, action 0
 54        # is the sole common safe action, so its group-average margin is q_bar[0].
 55        q_bar_feasible = q[:2].mean(0)
 56        robust_hinge = F.relu(margin - q_bar_feasible[0])
 57        loss = task + consistency_weight * consistency + safety_weight * robust_hinge
 58        opt.zero_grad(); loss.backward(); opt.step()
 59    with torch.no_grad():
 60        q = logits.softmax(-1)
 61        kappa = float((js(q[0], q[1]) + js(q[2], q[3])) / 2)
 62        feasible_violation = float(1 - q[:2, 0].mean())
 63        # Empty intersection has max_U min_x C(x,U)=0 exactly, independent of q.
 64        empty_robust_margin = 0.0
 65        empty_certificate = True
 66        return {
 67            "kappa": kappa, "feasible_violation": feasible_violation,
 68            "empty_robust_margin": empty_robust_margin,
 69            "empty_intersection_certificate": empty_certificate,
 70            "task_ce": float(F.cross_entropy(logits, labels)),
 71            "q": q.numpy().round(4).tolist(),
 72        }
 73
 74
 75def main():
 76    math = exact_math_check()
 77    structure = safety_structure_check()
 78    # First sweep the JS coefficient at zero safety weight: predicted kappa decline.
 79    consistency_weights = [0.0, 0.02, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0]
 80    consistency_sweep = {
 81        str(a): train(a, 0.0) for a in consistency_weights
 82    }
 83    # Then sweep robust-safety weight at fixed consistency: predicted feasible
 84    # violation decline, while the empty intersection remains certified impossible.
 85    safety_weights = [0.0, 0.02, 0.05, 0.1, 0.25, 0.5, 1.0, 2.0, 5.0]
 86    safety_sweep = {
 87        str(b): train(1.0, b) for b in safety_weights
 88    }
 89    baseline = consistency_sweep["0.0"]
 90    kappas = [consistency_sweep[str(a)]["kappa"] for a in consistency_weights]
 91    violations = [safety_sweep[str(b)]["feasible_violation"] for b in safety_weights]
 92    checks = {
 93        "P1_js_zero_is_exact": math["js_zero_equal"] == 0.0,
 94        "P2_kappa_reduced_high_consistency": kappas[-1] < kappas[0],
 95        "P2_monotone_kappa_fraction": float(np.mean(np.diff(kappas) <= 1e-5)),
 96        "P3_feasible_violation_reduced_high_safety": violations[-1] < violations[0],
 97        "P3_empty_intersection_always_certificate": all(
 98            safety_sweep[str(b)]["empty_intersection_certificate"] and
 99            safety_sweep[str(b)]["empty_robust_margin"] == 0.0 for b in safety_weights),
100    }
101    result = {
102        "seed": SEED, "math_check": math, "safety_structure": structure,
103        "baseline_task_only": baseline,
104        "consistency_sweep": consistency_sweep,
105        "safety_sweep": safety_sweep,
106        "prediction_checks": checks,
107        "prediction_summary": {
108            "P1": "JS is exactly zero for equal distributions and positive for unequal ones.",
109            "P2": "Increasing the JS coefficient reduces observational disagreement kappa.",
110            "P3": "Increasing robust-safety weight reduces feasible-group unsafe probability; empty intersections stay explicitly infeasible.",
111        },
112    }
113    print(json.dumps(result, indent=2))
114
115
116if __name__ == "__main__":
117    main()