cernpeerenv-zmumu-dqn

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

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.

from stable_baselines3 import DQN
from cern_hunt_env import CernHuntEnv

model = DQN.load("best_model_SingleAgent_dqn.zip")

env = CernHuntEnv(
    allow_pythia=True,
    infinite_data=True,
    reward_shaping=True,
    require_real_source=False,
)
obs, _ = env.reset(seed=42)
action, _ = model.predict(obs, deterministic=True)
obs, reward, terminated, truncated, info = env.step(int(action))
# info: label, best_inv_mass, mu_count
# terminated is always True
env.close()

model.predict is shown on one unbatched observation. For many events, loop or use SB3 evaluate_policy.

Published checkpoints

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.

Weight sweep on the mixed trio (aggregation="weighted", no retraining, 300 events, seed 777):

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

gymnasium>=0.29.0,<1.0.0
numpy>=1.25.0,<2.0.0
uproot>=5.0.0
vector>=0.9.0
stable-baselines3>=2.0.0
torch>=2.0.0
pythia8mc>=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.