Z→μμ event classifier (DQN/PPO trained) + peer-vote ensemble negative result
Task. Binary classification of simulated dimuon events: does this event contain an opposite-sign muon pair whose invariant mass falls in (80, 100) GeV?
Method. The classifier is trained with Stable-Baselines3 DQN or PPO inside a one-step Gymnasium wrapper. Each reset / step is one labeled event and then terminated=True. There is no horizon, no delayed credit assignment, and no control of a detector.
Label. Computed from the same padded muon kinematics the policy sees: at least two muons, opposite charge, invariant mass in (mass_min, mass_max) with defaults 80–100 GeV. This is not generator-level Z truth and not a CMS official tag.
Data for the reported numbers. Pythia8 Monte Carlo only (WeakSingleBoson:ffbar2gmZ, Z forced to μ⁺μ⁻, hat mass 60–120 GeV). The negative class is mostly off-window OS dimuons, not QCD background. None of the published checkpoints were trained or evaluated on real CMS ROOT Open Data.
If you want the classifier, load best_model_SingleAgent_dqn.zip at the repo root (not SingleAgent_DQN/best_model.zip). Peer-vote checkpoints are included as an ablation, not as a better model.
Scope
| Claim | Status |
|---|
| One-step binary classifier on muon kinematics | Yes |
| Label = OS pair in the 80–100 GeV mass window | Yes |
| SB3 DQN/PPO used as the optimizer | Yes |
| Sequential RL / POMDP / long horizon | No |
| Independent HEP truth label | No |
| Real CMS Open Data in these checkpoints | No |
| Z vs QCD tagging | No |
| Compared to a BDT or supervised MLP | No |
| Full mixed-ensemble weights published | No (mixed agent0 PPO is missing) |
What is included
cern_hunt_env.py — Gymnasium adapter (CernHuntEnv)
- Class labels exposed as
Discrete(2): 0 = no window pair, 1 = OS pair in (80, 100) GeV
- Fixed-shape
Dict observation: Muon_pt/eta/phi/mass/charge + validity mask, padded to max_muons=8
- Data order: optional CMS-style ROOT
Events tree → Pythia8 → Herwig CLI. Reported runs use Pythia8 only
- Training signal:
+1 / −1 for correct / incorrect class; optional mass-proximity bonus when the window label is 1 (reward_shaping=True in the reported runs)
- Each
step() classifies one event and terminates
peer_voting_env.py — inference-time vote wrapper (PeerVotingEnv)
- N independently trained classifiers see the same event and vote
- Aggregation:
majority, unanimous, or weighted
- Optional consensus bonus (
+0.1) when all votes agree
- Not multi-agent RL: agents do not interact over time
baseline_training.py — training and 300-event holdout eval
- Single-agent PPO and DQN; homogeneous 3×PPO votes; mixed 2×PPO+1×DQN votes
- Disjoint seed ranges (
STANDALONE_SEEDS, PEER_SEED_BASE)
- Reported training used
allow_pythia=True, require_real_source=False, reward_shaping=True, DummyVecEnv of 4 raw envs, MultiInputPolicy
weighted_voting_eval.py — inference-only weight sweep on the mixed trio
- Reloads trained members and sweeps the DQN vote weight
- Writes
results/weighted_voting_sweep.csv
NormalizedCernHuntEnv exists in the env file and was not used for the numbers below.
How to use
The Hugging Face repo is a flat checkout. Import the local modules, or install the package and use cernpeerenv.
1from stable_baselines3 import DQN
2from cern_hunt_env import CernHuntEnv
3
4model = DQN.load("best_model_SingleAgent_dqn.zip")
5
6env = CernHuntEnv(
7 allow_pythia=True,
8 infinite_data=True,
9 reward_shaping=True,
10 require_real_source=False,
11)
12obs, _ = env.reset(seed=42)
13action, _ = model.predict(obs, deterministic=True)
14obs, reward, terminated, truncated, info = env.step(int(action))
15# info: label, best_inv_mass, mu_count
16# terminated is always True
17env.close()
model.predict is shown on one unbatched observation. For many events, loop or use SB3 evaluate_policy.
Published checkpoints
| File | What it is |
|---|
best_model_SingleAgent_dqn.zip | Recommended. Single-agent DQN (92.7% on seed-888 holdout) |
best_model_SingleAgent_ppo.zip | Single-agent PPO (87.0% on seed-888 holdout) |
PeerVoting_3Agents_Mixed_agent1_PPO_106000_steps.zip | Mixed-ensemble member (PPO) |
PeerVoting_3Agents_Mixed_agent2_DQN_100000_steps.zip | Mixed-ensemble member (DQN) |
PeerVoting_3Agents_PPO_agent0_PPO_106000_steps.zip | One member of the homogeneous PPO trio |
Missing from this repo: mixed-ensemble agent0 PPO, and homogeneous-PPO agents 1 and 2. The mixed-vote table is not fully reproducible from published weights alone.
results/ holds learning-curve CSVs/PNGs and weighted_voting_sweep.csv.
Checkpoints are Stable-Baselines3 2.9.0, Gymnasium 0.29.1, PyTorch 2.x. Policy nets are small (policy.pth ≈ 65 KB); CPU inference is fine.
Validation results
All reported training used Pythia8 only. Accuracies count predicted class == window label.
| Method | Accuracy | Agreement | Eval |
|---|
| Random | 48.6% | n/a | 500 events, seed 0 |
| PPO (single) | 87.0% | n/a | 300 events, seed 888 |
| DQN (single) | 92.7% | n/a | 300 events, seed 888 |
| PPO majority vote (3 agents, homogeneous) | 87.7% | 97.0% | 300 events, seed 777 |
| Mixed vote, unweighted (2×PPO + 1×DQN) | 89.3% | 93.0% | 300 events, seed 777 |
Weight sweep on the mixed trio (aggregation="weighted", no retraining, 300 events, seed 777):
| DQN weight | Ensemble accuracy |
|---|
| 0.34 (≈uniform) | 89.3% |
| 0.40 | 89.3% |
| 0.50 | 93.7% |
| 0.60–0.90 | 93.7% (flat) |
Read these as one-seed, 300-event holdouts. A 1-point gap is sampling noise. Do not merge 92.7% and 93.7% into “DQN 92.7–93.7%”: they are different eval seeds, and 93.7% is a weighted vote that has already collapsed onto the DQN member.
Negative result. Once the DQN weight is ≥ 0.5, it outweighs the two PPO votes whenever they disagree, so the ensemble reproduces the DQN decision. That is not evidence that voting helps. The homogeneous PPO trio agreeing on 97% of events is the same story: seed-diverse copies collapse to nearly the same boundary.
Class balance is not logged. The random baseline near 50% only suggests the 80–100 GeV window does not make the Pythia sample extremely one-sided.
No supervised MLP, logistic regression, or BDT was run on the same features. A hard mass-window cut is the label; a linear model on the same 4-vectors would be the fair non-SB3 baseline and is not in this repo.
Dependencies
1gymnasium>=0.29.0,<1.0.0
2numpy>=1.25.0,<2.0.0
3uproot>=5.0.0
4vector>=0.9.0
5stable-baselines3>=2.0.0
6torch>=2.0.0
7pythia8mc>=1.0.0
Reported runs: Colab-class T4, ~15 minutes per 100k events. device="auto" picks CUDA when present. CPU training is practical at this network size.
Notes
Pythia8 only for every number in the tables (allow_pythia=True, require_real_source=False). The env can load a ROOT file with an Events tree via root_path=....
Single training seed per row. No multi-seed mean. 300-event eval sets are small. gamma=0.99 is the SB3 default even though each episode is one step.
This is a packaging baseline: Gymnasium plumbing, a window-label classifier, and a documented failure of naive vote aggregation.