dodox/NuCon
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
041e0ec1bd
NuCon/nucon/sim.py
Dominik Roth 65190dffea feat: uncertainty-aware training with penalty and abort 
sim.py:
- simulator.update(return_uncertainty=True) calls forward_with_uncertainty
  on kNN models and returns the GP std; returns None for NN or when not
  requested (no extra cost if unused)
- No state stored on simulator; caller decides what to do with the value

rl.py (NuconEnv and NuconGoalEnv):
- uncertainty_penalty_start: above this GP std, subtract a linear penalty
  from the reward (scaled by uncertainty_penalty_scale, default 1.0)
- uncertainty_abort: at or above this GP std, set truncated=True
- Only calls update(return_uncertainty=True) when either threshold is set
- Uncertainty only applies when using a simulator (kNN model); ignored otherwise

Example:
    simulator = NuconSimulator()
    simulator.load_model('reactor_knn.pkl')
    env = NuconGoalEnv(..., simulator=simulator,
                       uncertainty_penalty_start=0.3,
                       uncertainty_abort=0.7,
                       uncertainty_penalty_scale=2.0)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-12 18:37:09 +01:00

355 lines
16 KiB
Python
Raw Blame History

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 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):
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.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, return_uncertainty: bool = False):
"""Advance the simulator by time_step game-seconds.
If return_uncertainty=True and a kNN model is loaded, returns the GP
posterior std for this step (0 = on known data, ~1 = OOD).
Always returns None when using an NN model.
"""
uncertainty = self._update_reactor_state(time_step, return_uncertainty=return_uncertainty)
self.time += time_step
return uncertainty
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, return_uncertainty: bool = False):
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
uncertainty = None
if isinstance(self.model, ReactorDynamicsModel):
with torch.no_grad():
next_state = self.model.forward(state, time_step)
elif return_uncertainty:
next_state, uncertainty = self.model.forward_with_uncertainty(state, time_step)
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)
return uncertainty
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()
Reference in New Issue View Git Blame Copy Permalink