dodox/NuCon
1
1
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
f582e72151
NuCon/nucon/sim.py
Dominik Roth 0932bb353a feat: SAC+HER training on kNN-GP sim with direct bypass and scripts/ 
- nucon/rl.py: delta_action_scale action space, bool handling (>=0.5),
  direct sim read/write bypassing HTTP for ~2000fps env throughput;
  remove uncertainty_abort from training (use penalty-only), larger
  default batch sizes; fix _read_obs and step for in-process sim
- nucon/model.py: optimise _lookup with einsum squared-L2, vectorised
  rbf kernel; forward_with_uncertainty uses pre-built normalised arrays
- nucon/sim.py: _update_reactor_state writes outputs via setattr directly
- scripts/train_sac.py: moved from root; full SAC+HER example with kNN-GP
  sim, delta actions, uncertainty penalty, init_states
- scripts/collect_dataset.py: CLI tool to collect dynamics dataset from
  live game session (--steps, --delta, --out, --merge)
- README.md: add Scripts section, reference both scripts in training loop

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

359 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
params = self.parameters
state = {}
for param_id in self.model.input_params:
value = getattr(params, param_id, None)
if isinstance(value, Enum):
value = value.value
state[param_id] = 0.0 if value is None else 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)
# Write outputs directly — bypass sim.set() type-checking overhead
for param_id, value in next_state.items():
setattr(params, param_id, value)
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)
if value is None:
param = self._nucon[variable]
if param.enum_type is not None:
value = next(iter(param.enum_type)).value # first enum member's int value
else:
value = param.param_type() # int()->0, float()->0.0, bool()->False
if isinstance(value, Enum):
value = value.value
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