ENN Uncertainty-Quantification Module

Epistemic-uncertainty KPI for RL grid-operation recommendations, based on an Evidential Neural Network (ENN) trained by behavior cloning on the agent's own decisions. For each recommendation it returns two percentiles (0–100):

• epistemic_uncertainty_total_pctile — total epistemic uncertainty of the situation (the ENN vacuity u = K/S); high when the grid state is out-of-distribution for the agent.
• epistemic_uncertainty_action_pctile — epistemic uncertainty of the action the agent selected (variance of that action's predicted probability); None for a do-nothing action (not in the curated set).

Percentiles are reported (rather than raw values) because the raw measures sit in narrow bands; a percentile spreads them onto a readable 0–100 scale relative to a pre-computed reference distribution. No calibration step is needed on the consumer side — the reference is shipped (enn_pctile_calib.npz).

Environment

Python 3.9 or 3.10 is required. tensorflow==2.12.1, ray==2.5.1 and torch==2.1.2 publish no wheels for Python ≥ 3.11 (use a python:3.10-slim base image for containers — see Dockerfile). Grid2Op is pinned to 1.9.8, the same version used by the InteractiveAI simulator, so the action format stays consistent.

conda create -n enn_uq python=3.10 -y
conda activate enn_uq
pip install -r requirements.txt

Quick start — full reproduction

The scaler used by the module is created at ENN training time, so the complete reproduction of the experiment is: collect rollouts → train (the training exports the scaler, the metadata and the percentile calibration automatically) → run the example. For the CurriculumAgent nothing needs to be edited:

python training/collect_rollouts.py --agent curriculum --episodes 50 --out-dir data_curriculum
python training/train_enn.py --data-dir data_curriculum --out-dir models_curriculum --agent-name curriculum
python run_example.py

run_example.py is self-configuring: it locates the trained artifacts (newest scaler_params.json + enn_meta.json + calibration .npz, wherever the training wrote them), the curated action set, the ENN architecture module and the agent binaries, printing every resolved path (anything can be overridden in its CONFIG block). It then creates the l2rpn_icaps_2021_small environment (LightSim backend), loads the CurriculumAgent from curriculumagent/, rebuilds the training scaler from the exported JSON (no pickled scaler, so it reloads under any scikit-learn version), loads the calibration, and calls assess_recommendation on live observations, printing the two percentiles per step and the recommendations list in the InteractiveAI format. If no trained artifacts exist yet, it prints the exact training commands to run first.

Artifacts

File	What it is
recommendation_uncertainty.py	the module: load_calibration, assess_recommendation
src/enn_models.py	ENN architecture (EvidentialNetwork)
data_<agent>/actions.npy (produced by collection)	curated action set — rows are action.to_vect(), deduplicated from the rollouts
models_<agent>/ (produced by training)	enn_<agent>.pth, scaler_params.json (scaler mean/std as JSON — the scaler is created at training time), enn_meta.json, enn_pctile_calib.npz
curriculumagent/	agent binaries (L2RPN submission layout)

Usage in three lines

from recommendation_uncertainty import load_calibration, assess_recommendation

calibration = load_calibration("models_curriculum/enn_pctile_calib.npz",
                               scaler=scaler,
                               action_set="data_curriculum/actions.npy",
                               class_mapping="models_curriculum/enn_meta.json")
info = assess_recommendation(obs, agent, enn, calibration)
# {"chosen_action_id": ..., "epistemic_uncertainty_total_pctile": ...,
#  "epistemic_uncertainty_action_pctile": ...}

The agent returns a Grid2Op action object (not an index): the module locates it in the curated action set to obtain its index, then maps it to the ENN label via class_mapping. A do-nothing action is not in the set, so its per-action value is None (the situation-level value is still produced).

InteractiveAI output format

Each recommendation is a dictionary in the recommendations list; the two percentiles go into the kpis field alongside efficiency_of_the_reco:

{
  "title": "Topological recommendation (CurriculumAgent)",
  "description": "...",
  "use_case": "PowerGrid",
  "agent_type": 2,
  "actions": [ { "...": "grid2op serializable action dict" } ],
  "kpis": {
    "efficiency_of_the_reco": null,
    "epistemic_uncertainty_total_pctile": 47.8,
    "epistemic_uncertainty_action_pctile": 3.1
  }
}

See to_interactiveai() in run_example.py for the exact mapping.

Tests (no trained weights needed)

python tests/validate_module.py

Validates the module end-to-end with a synthetic ENN exposing the same forward interface: action location in the curated set, do-nothing handling, percentile-mapping monotonicity and bounds, and InteractiveAI serialisation. Useful to confirm an environment is correctly set up before touching the real binaries.

Training the ENN for a new agent

The module is agent-agnostic — any object exposing agent.act(obs, reward, done) -> grid2op action can be wrapped. The same two scripts used above handle any agent: --agent expert in training/collect_rollouts.py (after plugging in the ExpertAgent constructor) collects rollouts and curates the action set automatically by deduplication; training/train_enn.py then trains, exports the scaler and metadata, and calibrates. Full step-by-step guide: TRAINING.md.

Docker

docker build -t enn-uq .
docker run --rm enn-uq            # runs run_example.py as a reproducibility check

Maintainers

export_artifacts.py regenerates the JSON artifacts from an existing fitted scaler (.pkl) and checkpoint, for the case where a training was done outside training/train_enn.py. It is fully self-configuring (it infers input_dim/num_classes from the checkpoint's state_dict and cross-checks everything). With the standard training flow it is not needed — train_enn.py already exports everything.

License

Mozilla Public License 2.0 — see LICENSE.