Dominik Roth
6cb93ad56d
NuconSimulator now accepts uncertainty_threshold (default None = disabled).
When set and using a kNN model, _update_reactor_state() calls
forward_with_uncertainty() and raises HighUncertaintyError if the GP
posterior std exceeds the threshold.
NuconEnv and NuconGoalEnv catch HighUncertaintyError in step() and
return truncated=True, so SB3 bootstraps the value rather than treating
OOD regions as terminal states.
Usage:
simulator = NuconSimulator(uncertainty_threshold=0.3)
# episodes are cut short when the policy wanders OOD
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
361 lines
17 KiB
Python
361 lines
17 KiB
Python
import random
|
|
from typing import Dict, Union, Any, Tuple, List
|
|
from enum import Enum
|
|
from flask import Flask, request, jsonify
|
|
from nucon import Nucon, ParameterEnum, PumpStatus, PumpDryStatus, PumpOverloadStatus, BreakerStatus
|
|
import threading
|
|
import torch
|
|
from nucon.model import ReactorDynamicsModel, ReactorKNNModel
|
|
import pickle
|
|
|
|
class HighUncertaintyError(Exception):
|
|
"""Raised by NuconSimulator when the dynamics model uncertainty exceeds the threshold.
|
|
|
|
Caught by NuconEnv/NuconGoalEnv and returned as truncated=True so the RL
|
|
algorithm bootstraps the value rather than treating it as a terminal state.
|
|
"""
|
|
def __init__(self, uncertainty: float, threshold: float):
|
|
self.uncertainty = uncertainty
|
|
self.threshold = threshold
|
|
super().__init__(f"Model uncertainty {uncertainty:.3f} exceeded threshold {threshold:.3f}")
|
|
|
|
|
|
class OperatingState(Enum):
|
|
# Tuple indicates a range of values, while list indicates a set of possible values
|
|
OFFLINE = {
|
|
'CORE_TEMP': (18.0, 22.0),
|
|
'CORE_TEMP_OPERATIVE': [306.0],
|
|
'CORE_TEMP_MAX': [1000.0],
|
|
'CORE_TEMP_MIN': [0.0],
|
|
'CORE_TEMP_RESIDUAL': [False],
|
|
'CORE_PRESSURE': (0.9, 1.1),
|
|
'CORE_PRESSURE_MAX': [1.5],
|
|
'CORE_PRESSURE_OPERATIVE': [1.0],
|
|
'CORE_INTEGRITY': [100.0],
|
|
'CORE_WEAR': [0.0],
|
|
'CORE_STATE': ['OFFLINE'],
|
|
'CORE_STATE_CRITICALITY': [0.0],
|
|
'CORE_CRITICAL_MASS_REACHED': [False],
|
|
'CORE_CRITICAL_MASS_REACHED_COUNTER': [0],
|
|
'CORE_IMMINENT_FUSION': [False],
|
|
'CORE_READY_FOR_START': [False],
|
|
'CORE_STEAM_PRESENT': [False],
|
|
'CORE_HIGH_STEAM_PRESENT': [False],
|
|
'TIME': ['00:00:00'],
|
|
'TIME_STAMP': ['1970-01-01 00:00:00'],
|
|
'COOLANT_CORE_STATE': ['INACTIVE'],
|
|
'COOLANT_CORE_PRESSURE': (0.9, 1.1),
|
|
'COOLANT_CORE_MAX_PRESSURE': [1.5],
|
|
'COOLANT_CORE_VESSEL_TEMPERATURE': (18.0, 22.0),
|
|
'COOLANT_CORE_QUANTITY_IN_VESSEL': [0.0],
|
|
'COOLANT_CORE_PRIMARY_LOOP_LEVEL': [0.0],
|
|
'COOLANT_CORE_FLOW_SPEED': [0.0],
|
|
'COOLANT_CORE_FLOW_ORDERED_SPEED': [0.0],
|
|
'COOLANT_CORE_FLOW_REACHED_SPEED': [False],
|
|
'COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT': [3],
|
|
'COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT': [2],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_STATUS': [PumpStatus.INACTIVE],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_STATUS': [PumpStatus.INACTIVE],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_STATUS': [PumpStatus.INACTIVE],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_ORDERED_SPEED': [0.0],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_ORDERED_SPEED': [0.0],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_ORDERED_SPEED': [0.0],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_SPEED': [0.0],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_SPEED': [0.0],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_SPEED': [0.0],
|
|
'RODS_STATUS': ['INACTIVE'],
|
|
'GENERATOR_0_KW': [0.0],
|
|
'GENERATOR_1_KW': [0.0],
|
|
'GENERATOR_2_KW': [0.0],
|
|
'GENERATOR_0_V': [0.0],
|
|
'GENERATOR_1_V': [0.0],
|
|
'GENERATOR_2_V': [0.0],
|
|
'GENERATOR_0_A': [0.0],
|
|
'GENERATOR_1_A': [0.0],
|
|
'GENERATOR_2_A': [0.0],
|
|
'GENERATOR_0_HERTZ': [0.0],
|
|
'GENERATOR_1_HERTZ': [0.0],
|
|
'GENERATOR_2_HERTZ': [0.0],
|
|
'GENERATOR_0_BREAKER': [BreakerStatus.OPEN],
|
|
'GENERATOR_1_BREAKER': [BreakerStatus.OPEN],
|
|
'GENERATOR_2_BREAKER': [BreakerStatus.OPEN],
|
|
'STEAM_TURBINE_0_RPM': [0.0],
|
|
'STEAM_TURBINE_1_RPM': [0.0],
|
|
'STEAM_TURBINE_2_RPM': [0.0],
|
|
'STEAM_TURBINE_0_TEMPERATURE': (18.0, 22.0),
|
|
'STEAM_TURBINE_1_TEMPERATURE': (18.0, 22.0),
|
|
'STEAM_TURBINE_2_TEMPERATURE': (18.0, 22.0),
|
|
'STEAM_TURBINE_0_PRESSURE': (0.9, 1.1),
|
|
'STEAM_TURBINE_1_PRESSURE': (0.9, 1.1),
|
|
'STEAM_TURBINE_2_PRESSURE': (0.9, 1.1),
|
|
}
|
|
NOMINAL = {
|
|
'CORE_TEMP': (290.0, 370.0),
|
|
'CORE_TEMP_OPERATIVE': [306.0],
|
|
'CORE_TEMP_MAX': [1000.0],
|
|
'CORE_TEMP_MIN': [0.0],
|
|
'CORE_TEMP_RESIDUAL': [False],
|
|
'CORE_PRESSURE': (14.5, 15.5),
|
|
'CORE_PRESSURE_MAX': [16.0],
|
|
'CORE_PRESSURE_OPERATIVE': [15.0],
|
|
'CORE_INTEGRITY': (99.0, 100.0),
|
|
'CORE_WEAR': (0.0, 1.0),
|
|
'CORE_STATE': ['ACTIVE'],
|
|
'CORE_STATE_CRITICALITY': (0.9, 1.1),
|
|
'CORE_CRITICAL_MASS_REACHED': [True],
|
|
'CORE_CRITICAL_MASS_REACHED_COUNTER': [1],
|
|
'CORE_IMMINENT_FUSION': [False],
|
|
'CORE_READY_FOR_START': [True],
|
|
'CORE_STEAM_PRESENT': [True],
|
|
'CORE_HIGH_STEAM_PRESENT': [False],
|
|
'TIME': ['00:00:00'],
|
|
'TIME_STAMP': ['1970-01-01 00:00:00'],
|
|
'COOLANT_CORE_STATE': ['ACTIVE'],
|
|
'COOLANT_CORE_PRESSURE': (13.5, 14.5),
|
|
'COOLANT_CORE_MAX_PRESSURE': [15.0],
|
|
'COOLANT_CORE_VESSEL_TEMPERATURE': (270.0, 290.0),
|
|
'COOLANT_CORE_QUANTITY_IN_VESSEL': (95.0, 100.0),
|
|
'COOLANT_CORE_PRIMARY_LOOP_LEVEL': (95.0, 100.0),
|
|
'COOLANT_CORE_FLOW_SPEED': (9.5, 10.5),
|
|
'COOLANT_CORE_FLOW_ORDERED_SPEED': [10.0],
|
|
'COOLANT_CORE_FLOW_REACHED_SPEED': [True],
|
|
'COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT': [3],
|
|
'COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT': [3],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_STATUS': [PumpStatus.ACTIVE_SPEED_REACHED],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_STATUS': [PumpStatus.ACTIVE_SPEED_REACHED],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_STATUS': [PumpStatus.ACTIVE_SPEED_REACHED],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_DRY_STATUS': [PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_OVERLOAD_STATUS': [PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD],
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_ORDERED_SPEED': (95.0, 100.0),
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_ORDERED_SPEED': (95.0, 100.0),
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_ORDERED_SPEED': (95.0, 100.0),
|
|
'COOLANT_CORE_CIRCULATION_PUMP_0_SPEED': (95.0, 100.0),
|
|
'COOLANT_CORE_CIRCULATION_PUMP_1_SPEED': (95.0, 100.0),
|
|
'COOLANT_CORE_CIRCULATION_PUMP_2_SPEED': (95.0, 100.0),
|
|
'RODS_STATUS': ['ACTIVE'],
|
|
'GENERATOR_0_KW': (950.0, 1050.0),
|
|
'GENERATOR_1_KW': (950.0, 1050.0),
|
|
'GENERATOR_2_KW': (950.0, 1050.0),
|
|
'GENERATOR_0_V': (380.0, 420.0),
|
|
'GENERATOR_1_V': (380.0, 420.0),
|
|
'GENERATOR_2_V': (380.0, 420.0),
|
|
'GENERATOR_0_A': (2375.0, 2625.0),
|
|
'GENERATOR_1_A': (2375.0, 2625.0),
|
|
'GENERATOR_2_A': (2375.0, 2625.0),
|
|
'GENERATOR_0_HERTZ': (49.5, 50.5),
|
|
'GENERATOR_1_HERTZ': (49.5, 50.5),
|
|
'GENERATOR_2_HERTZ': (49.5, 50.5),
|
|
'GENERATOR_0_BREAKER': [BreakerStatus.CLOSED],
|
|
'GENERATOR_1_BREAKER': [BreakerStatus.CLOSED],
|
|
'GENERATOR_2_BREAKER': [BreakerStatus.CLOSED],
|
|
'STEAM_TURBINE_0_RPM': (2950.0, 3050.0),
|
|
'STEAM_TURBINE_1_RPM': (2950.0, 3050.0),
|
|
'STEAM_TURBINE_2_RPM': (2950.0, 3050.0),
|
|
'STEAM_TURBINE_0_TEMPERATURE': (270.0, 290.0),
|
|
'STEAM_TURBINE_1_TEMPERATURE': (270.0, 290.0),
|
|
'STEAM_TURBINE_2_TEMPERATURE': (270.0, 290.0),
|
|
'STEAM_TURBINE_0_PRESSURE': (13.5, 14.5),
|
|
'STEAM_TURBINE_1_PRESSURE': (13.5, 14.5),
|
|
'STEAM_TURBINE_2_PRESSURE': (13.5, 14.5),
|
|
}
|
|
|
|
class NuconSimulator:
|
|
class Parameters:
|
|
def __init__(self, nucon: Nucon):
|
|
for param_name in nucon.get_all_readable():
|
|
setattr(self, param_name, None)
|
|
|
|
def __init__(self, host: str = 'localhost', port: int = 8786, uncertainty_threshold: float = None):
|
|
self._nucon = Nucon()
|
|
self.parameters = self.Parameters(self._nucon)
|
|
self.host = host
|
|
self.port = port
|
|
self.time = 0.0
|
|
self.allow_all_writes = False
|
|
self.uncertainty_threshold = uncertainty_threshold
|
|
self.set_state(OperatingState.OFFLINE)
|
|
self.model = None
|
|
self.readable_params = list(self._nucon.get_all_readable().keys())
|
|
self.non_writable_params = [name for name, param in self._nucon.get_all_readable().items() if not param.is_writable]
|
|
self._run(host, port)
|
|
|
|
def get(self, parameter: Union[str, Any]) -> Any:
|
|
if isinstance(parameter, str):
|
|
return getattr(self.parameters, parameter)
|
|
return getattr(self.parameters, parameter.id)
|
|
|
|
def set(self, parameter: Union[str, Any], value: Any, force: bool = False) -> None:
|
|
if isinstance(parameter, str):
|
|
param_obj = self._nucon[parameter]
|
|
else:
|
|
param_obj = parameter
|
|
|
|
if not param_obj.is_writable and not force and not self.allow_all_writes:
|
|
raise ValueError(f"Parameter {param_obj.id} is not writable")
|
|
|
|
# Convert value to the correct type
|
|
try:
|
|
if param_obj.enum_type:
|
|
if isinstance(value, str):
|
|
value = param_obj.enum_type[value.upper()]
|
|
elif isinstance(value, int):
|
|
value = param_obj.enum_type(value)
|
|
elif not isinstance(value, param_obj.enum_type):
|
|
raise ValueError(f"Invalid enum value for {param_obj.id}")
|
|
else:
|
|
value = param_obj.param_type(value)
|
|
except (ValueError, KeyError):
|
|
raise ValueError(f"Invalid type for parameter {param_obj.id}. Expected {param_obj.param_type}")
|
|
|
|
# Check range if not forced
|
|
if not force and param_obj.min_val is not None and param_obj.max_val is not None:
|
|
if not param_obj.min_val <= value <= param_obj.max_val:
|
|
raise ValueError(f"Value {value} is out of range for parameter {param_obj.id}. "
|
|
f"Valid range: [{param_obj.min_val}, {param_obj.max_val}]")
|
|
|
|
setattr(self.parameters, param_obj.id, value)
|
|
|
|
def set_allow_all_writes(self, allow: bool) -> None:
|
|
self.allow_all_writes = allow
|
|
|
|
def update(self, time_step: float) -> None:
|
|
self._update_reactor_state(time_step)
|
|
self.time += time_step
|
|
|
|
def set_model(self, model) -> None:
|
|
"""Set a pre-loaded ReactorDynamicsModel or ReactorKNNModel directly."""
|
|
self.model = model
|
|
if isinstance(model, ReactorDynamicsModel):
|
|
self.model.eval()
|
|
|
|
def load_model(self, model_path: str) -> None:
|
|
"""Load a model from a file. .pkl → ReactorKNNModel, otherwise → ReactorDynamicsModel (torch)."""
|
|
try:
|
|
if model_path.endswith('.pkl'):
|
|
with open(model_path, 'rb') as f:
|
|
self.model = pickle.load(f)
|
|
print(f"kNN model loaded from {model_path}")
|
|
else:
|
|
# Reconstruct shell from the saved state dict; input/output params
|
|
# are stored inside the checkpoint.
|
|
checkpoint = torch.load(model_path, weights_only=False)
|
|
if isinstance(checkpoint, dict) and 'input_params' in checkpoint:
|
|
self.model = ReactorDynamicsModel(checkpoint['input_params'], checkpoint['output_params'])
|
|
self.model.load_state_dict(checkpoint['state_dict'])
|
|
else:
|
|
# Legacy: plain state dict — fall back using sim readable/non-writable lists
|
|
self.model = ReactorDynamicsModel(self.readable_params, self.non_writable_params)
|
|
self.model.load_state_dict(checkpoint)
|
|
self.model.eval()
|
|
print(f"NN model loaded from {model_path}")
|
|
except Exception as e:
|
|
print(f"Error loading model: {str(e)}")
|
|
self.model = None
|
|
|
|
def _update_reactor_state(self, time_step: float) -> None:
|
|
if not self.model:
|
|
raise ValueError("Model not set. Please load a model using load_model() or set_model().")
|
|
|
|
# Build state dict using only the params the model knows about
|
|
state = {}
|
|
for param_id in self.model.input_params:
|
|
value = getattr(self.parameters, param_id, None)
|
|
if isinstance(value, Enum):
|
|
value = value.value
|
|
if value is None:
|
|
value = 0.0 # fallback for params not initialised in sim state
|
|
state[param_id] = value
|
|
|
|
# Forward pass — same interface for both NN and kNN
|
|
if isinstance(self.model, ReactorDynamicsModel):
|
|
with torch.no_grad():
|
|
next_state = self.model.forward(state, time_step)
|
|
else:
|
|
if self.uncertainty_threshold is not None:
|
|
next_state, uncertainty = self.model.forward_with_uncertainty(state, time_step)
|
|
if uncertainty > self.uncertainty_threshold:
|
|
raise HighUncertaintyError(uncertainty, self.uncertainty_threshold)
|
|
else:
|
|
next_state = self.model.forward(state, time_step)
|
|
|
|
# Update only the output params the model predicts
|
|
for param_id, value in next_state.items():
|
|
try:
|
|
self.set(param_id, value, force=True)
|
|
except (ValueError, KeyError):
|
|
pass # ignore params that can't be set (type mismatch, unknown)
|
|
|
|
def set_state(self, state: OperatingState) -> None:
|
|
self._sample_parameters_from_state(state)
|
|
|
|
def _sample_parameters_from_state(self, state: OperatingState) -> None:
|
|
for param_name, value_spec in state.value.items():
|
|
param = self._nucon.get_all_readable()[param_name]
|
|
if isinstance(value_spec, tuple):
|
|
value = random.uniform(*value_spec)
|
|
elif isinstance(value_spec, list):
|
|
value = random.choice(value_spec)
|
|
else:
|
|
raise ValueError(f"Invalid value specification for parameter {param_name}")
|
|
self.set(param, value, force=True)
|
|
|
|
def _run(self, host: str = 'localhost', port: int = 8786, debug: bool = False):
|
|
app = Flask(__name__)
|
|
|
|
@app.route('/', methods=['GET'])
|
|
def get_parameter():
|
|
variable = request.args.get('variable', '').upper()
|
|
if not variable:
|
|
return jsonify({"error": "No variable specified"}), 400
|
|
|
|
if variable == 'WEBSERVER_LIST_VARIABLES':
|
|
return '\n'.join(self._nucon._parameters.keys()), 200
|
|
|
|
if variable == 'WEBSERVER_BATCH_GET':
|
|
value_arg = request.args.get('value', '')
|
|
names = [n.strip().upper() for n in value_arg.split(',') if n.strip()]
|
|
if not names:
|
|
names = list(self._nucon._parameters.keys())
|
|
result = {}
|
|
for name in names:
|
|
try:
|
|
result[name] = str(self.get(name))
|
|
except (KeyError, AttributeError):
|
|
pass
|
|
return jsonify(result), 200
|
|
|
|
try:
|
|
value = self.get(variable)
|
|
return str(value), 200
|
|
except (KeyError, AttributeError):
|
|
return jsonify({"error": f"Unknown variable: {variable}"}), 404
|
|
|
|
@app.route('/', methods=['POST'])
|
|
def set_parameter():
|
|
variable = request.args.get('variable')
|
|
value = request.args.get('value')
|
|
if not variable or value is None:
|
|
return jsonify({"error": "Both variable and value must be specified"}), 400
|
|
|
|
try:
|
|
self.set(variable, value)
|
|
return "OK", 200
|
|
except ValueError as e:
|
|
return jsonify({"error": str(e)}), 400
|
|
except KeyError:
|
|
return jsonify({"error": f"Unknown variable: {variable}"}), 404
|
|
|
|
def run_simulator(host='localhost', port=8786, debug=False):
|
|
app.run(host=host, port=port, debug=debug)
|
|
|
|
threading.Thread(target=run_simulator, args=(host, port, debug), daemon=True).start() |