1import json, math, random
  2from dataclasses import dataclass
  3from pathlib import Path
  4import numpy as np
  5
  6@dataclass(frozen=True)
  7class Transition:
  8    name: str
  9    inp: int
 10    out: int
 11
 12class PetriNet:
 13    def __init__(self, n_places, transitions, initial, admissible):
 14        self.n_places = n_places
 15        self.transitions = transitions
 16        self.initial = tuple(initial)
 17        self.admissible = {tuple(m) for m in admissible}
 18
 19    def enabled(self, t, m):
 20        return m[t.inp] >= 1
 21
 22    def fire(self, t, m):
 23        x = list(m); x[t.inp] -= 1; x[t.out] += 1
 24        return tuple(x)
 25
 26    def allowed(self, m):
 27        return [t for t in self.transitions
 28                if self.enabled(t, m) and self.fire(t, m) in self.admissible]
 29
 30
 31def make_rover():
 32    # One-hot marking: position x battery b, plus one unsafe sink.
 33    states = [(x, b) for x in range(5) for b in range(3)]
 34    place = {s:i for i, s in enumerate(states)}
 35    sink = len(states)
 36    actions = ["left", "right", "recharge", "wait"]
 37    transitions = []
 38    for s in states:
 39        x, b = s
 40        for a in actions:
 41            valid = True
 42            if a == "left":
 43                valid = x > 0 and b > 0
 44                ns = (x-1, b-1) if valid else None
 45            elif a == "right":
 46                valid = x < 4 and b > 0
 47                ns = (x+1, b-1) if valid else None
 48            elif a == "recharge":
 49                valid = x == 0 and b < 2
 50                ns = (x, 2) if valid else None
 51            else:
 52                ns = s
 53            transitions.append(Transition(f"{s}:{a}", place[s], place[ns] if valid else sink))
 54    init = tuple(1 if i == place[(0, 2)] else 0 for i in range(sink+1))
 55    # M_obs is the exact reachable set under safe transitions, not merely all bit patterns.
 56    seen = {init}; frontier = [init]
 57    temp = PetriNet(sink+1, transitions, init, seen)
 58    while frontier:
 59        m = frontier.pop()
 60        for t in temp.transitions:
 61            if temp.enabled(t, m):
 62                q = temp.fire(t, m)
 63                if q[sink] == 0 and q not in seen:
 64                    seen.add(q); frontier.append(q)
 65    return PetriNet(sink+1, transitions, init, seen), states, actions, place, sink
 66
 67
 68def argmax_action(logits):
 69    return int(np.argmax(logits))
 70
 71
 72def run_policy(net, states, actions, place, sink, shield, rng, episodes=10000):
 73    # Fixed random linear neural policy: same scores are used by both controllers.
 74    d = len(states)
 75    W = rng.normal(0, 1, size=(len(actions), 2))
 76    violations = rejects = deadlocks = steps = successes = 0
 77    for _ in range(episodes):
 78        m = net.initial
 79        for _step in range(40):
 80            idx = int(np.argmax(np.asarray(m[:-1])))
 81            x, b = states[idx]
 82            logits = W @ np.array([x / 4.0, b / 2.0])
 83            candidates = [t for t in net.transitions if t.inp == idx]
 84            # actions are ordered consistently by construction.
 85            if shield:
 86                allowed = [t for t in candidates if net.fire(t, m) in net.admissible]
 87                if not allowed:
 88                    deadlocks += 1; break
 89                best = int(np.argmax([logits[actions.index(t.name.split(':')[1])] for t in allowed]))
 90                t = allowed[best]
 91                if int(np.argmax(logits)) != actions.index(t.name.split(':')[1]): rejects += 1
 92            else:
 93                ai = argmax_action(logits); t = candidates[ai]
 94                if net.fire(t, m)[sink] == 1:
 95                    violations += 1; break
 96            m = net.fire(t, m); steps += 1
 97            # Reaching x=4 is a simple task success signal.
 98            if states[int(np.argmax(np.asarray(m[:-1])))][0] == 4:
 99                successes += 1; break
100    return {"violations": violations, "rejections": rejects, "deadlocks": deadlocks,
101            "steps": steps, "successes": successes, "episodes": episodes}
102
103
104def rejection_sweep():
105    # Prediction: with equal logits and uniform sampling, rejection probability = 1-k/n.
106    rows = []
107    rng = np.random.default_rng(123)
108    for n in [2, 4, 8, 16]:
109        for k in range(1, n):
110            trials = 20000
111            chosen = rng.integers(0, n, size=trials)
112            observed = float(np.mean(chosen >= k))
113            predicted = 1.0 - k / n
114            rows.append({"n": n, "allowed": k, "predicted": predicted, "observed": observed,
115                         "abs_error": abs(observed-predicted)})
116    return rows
117
118
119def main():
120    random.seed(7); np.random.seed(7)
121    net, states, actions, place, sink = make_rover()
122    # Core math checks: every accepted successor is in M_obs; BFS is finite and nonempty.
123    closure_bad = sum(net.fire(t, m) not in net.admissible
124                      for m in net.admissible for t in net.allowed(m))
125    # Compare exact same policy logits with and without the transition mask.
126    base = run_policy(net, states, actions, place, sink, False, np.random.default_rng(99))
127    shield = run_policy(net, states, actions, place, sink, True, np.random.default_rng(99))
128    sweep = rejection_sweep()
129    max_err = max(r["abs_error"] for r in sweep)
130    result = {
131      "petri_net": {"places": net.n_places, "transitions": len(net.transitions),
132                    "admissible_markings": len(net.admissible), "closure_bad": closure_bad},
133      "predictions": {
134        "inductive_closure": {"predicted": "0 unsafe accepted successors", "observed": closure_bad},
135        "uniform_rejection": {"predicted": "P(reject)=1-k/n", "max_abs_error": max_err,
136                               "rows": sweep},
137        "shield_violations": {"predicted": 0, "observed": shield["violations"]}
138      },
139      "controller_comparison": {"baseline_unshielded": base, "idea_shielded": shield}
140    }
141    Path("results.json").write_text(json.dumps(result, indent=2))
142    print(json.dumps(result, indent=2))
143
144if __name__ == "__main__":
145    main()