fix: don't override primary pump speed; warn in TUI if far from suggested 65%

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

README.md

@@ -1,22 +1,28 @@
-# NuCon (Nucleares Controller)
+<div align="center">
+  <img src='./logo.svg' width="250px">
+  <h2>NuCon</h2>
+  <br>
+</div>
 
-NuCon is a Python library designed to interface with and control parameters in Nucleares, a nuclear reactor simulation game. It provides a robust, type-safe foundation for reading and writing game parameters, allowing users to easily create their own automations and control systems.
+NuCon (Nucleares Controller) is a Python library designed to interface with and control parameters in [Nucleares](https://store.steampowered.com/app/1428420/Nucleares/), a nuclear reactor simulation game. It provides a robust, type-safe foundation for reading and writing game parameters, allowing users to easily create their own automations and control systems.
 
-NuCon further provides a work in progress implementation of a reinforcement learning environment for training control policies.
+NuCon further provides a reinforcement learning environment for training control policies and a simulator based on model learning.
 
 ## Features
 
 - Enum-based parameter system for type safety and code clarity
 - Support for various parameter types including floats, integers, booleans, strings, and custom enums
 - Read and write capabilities for game parameters
-- Custom truthy values for status enums to simplify conditional logic
+- Reinforcement learning environment for training control policies
-- Dummy mode for testing without connecting to the game
+- Built-in simulator for rapid prototyping and testing
-- Batch operations for getting multiple parameters at once
+- Model learning for dynamics prediction
 
 ## Installation
 
 To install NuCon, clone this repository and install via pip:
 ```bash
+git clone https://git.dominik-roth.eu/dodox/NuCon
+cd NuCon
 pip install -e .
 ```
 
@@ -25,92 +31,119 @@ pip install -e .
 Here's a basic example of how to use NuCon:
 
 ```python
-from nucon import Nucon, BreakerStatus
+from nucon import Nucon
 
-# Set the base URL for the game's API (if different from default)
+nucon = Nucon()
-Nucon.set_base_url("http://localhost:8080/")
+# or nucon = Nucon(host='localhost', port=8786)
 
 # Enable dummy mode for testing (optional)
-Nucon.set_dummy_mode(True)
+nucon.set_dummy_mode(True)
 
 # Read a parameter
-core_temp = Nucon.CORE_TEMP.value
+core_temp = nucon.CORE_TEMP.value
 print(f"Core Temperature: {core_temp}")
+# >> Core Temperature: 500.0
 
 # Read a parameter with an enum type
-pump_status = Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value
+pump_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value
 print(f"Pump 0 Status: {pump_status}")
+if nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value:
+    print('Pump 0 is active.')
+# >> Pump 0 Status: PumpStatus.INACTIVE
 
-# Write to a parameter
+# Write to a parameter (has no effect in dummy mode)
-Nucon.GENERATOR_0_BREAKER.value = BreakerStatus.OPEN # or True
+nucon.RODS_POS_ORDERED.value = 50
-print(f"Generator 0 Breaker Status: {Nucon.GENERATOR_0_BREAKER.value}")
+print(f"Rods Position Ordered: {nucon.RODS_POS_ORDERED.value}")
+# >> Rods Position Ordered: 50.0
 
-# Use custom truthy values
+# The repr of an attribute contains all contained info
-if Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value:
+nucon.CORE_TEMP
-    print("Pump 0 is active")
+# >> NuconParameter(id='CORE_TEMP', value=500.0, param_type=float, is_writable=False)
 ```
 
 ## API Reference
 
-The `Nucon` enum contains all available parameters. Each parameter is defined with:
+The `nucon` instance contains all available parameters.
-- An ID (string)
-- A type (float, int, bool, str, or a custom Enum)
-- A boolean indicating whether it's writable
 
 Parameter properties:
-- `Nucon.<PARAMETER>.value`: Get or set the current value of the parameter. Assigning a new value will write it to the game.
+- `nucon.<PARAMETER>.value`: Get or set the current value of the parameter. Assigning a new value will write it to the game.
-- `Nucon.<PARAMETER>.param_type`: Get the type of the parameter
+- `nucon.<PARAMETER>.param_type`: Get the type of the parameter
-- `Nucon.<PARAMETER>.is_writable`: Check if the parameter is writable
+- `nucon.<PARAMETER>.is_writable`: Check if the parameter is writable
-- `Nucon.<PARAMETER>.enum_type`: Get the enum type of the parameter if it's an enum, otherwise None
+- `nucon.<PARAMETER>.is_readable`: `False` for write-only parameters (e.g. VALVE_OPEN, CORE_SCRAM_BUTTON). Reading raises `AttributeError`.
+- `nucon.<PARAMETER>.is_cheat`: `True` for game-event triggers (all `FUN_*`). Writing raises `ValueError` unless `cheat_mode=True`.
+- `nucon.<PARAMETER>.enum_type`: Get the enum type of the parameter if it's an enum, otherwise None
+- `nucon.<PARAMETER>.unit`: Unit string if defined (e.g. `'°C'`, `'bar'`, `'%'`)
 
 Parameter methods:
-- `Nucon.<PARAMETER>.read()`: Get the current value of the parameter (alias for `value`)
+- `nucon.<PARAMETER>.read()`: Get the current value of the parameter (alias for `value`)
-- `Nucon.<PARAMETER>.write(new_value, force=False)`: Write a new value to the parameter. `force` will try to write even if the parameter is known as non-writable.
+- `nucon.<PARAMETER>.write(new_value, force=False)`: Write a new value to the parameter. `force` will try to write even if the parameter is known as non-writable or out of known allowed range.
 
 Class methods:
-- `Nucon.get(parameter)`: Get the value of a specific parameter. Also accepts string parameter names.
+- `nucon.get(parameter)`: Get the value of a specific parameter. Also accepts string parameter names.
-- `Nucon.set(parameter, value, force=False)`: Set the value of a specific parameter. Also accepts string parameter names.
+- `nucon.set(parameter, value, force=False)`: Set the value of a specific parameter. Also accepts string parameter names. `force` bypasses writable/range/cheat checks.
-- `Nucon.get_all_readable()`: Get a list of all readable parameters (which is all parameters)
+- `nucon.get_all_readable()`: Get a dict of all readable parameters.
-- `Nucon.get_all_writable()`: Get a list of all writable parameters
+- `nucon.get_all_writable()`: Get a dict of all writable parameters (includes write-only params).
-- `Nucon.get_all()`: Get all parameter values as a dictionary
+- `nucon.get_all()`: Get all readable parameter values as a dictionary.
-- `Nucon.get_multiple(params)`: Get values for multiple specified parameters
+- `nucon.get_all_iter()`: Get all readable parameter values as a generator.
-- `Nucon.set_base_url(url)`: Set the base URL for the game's API
+- `nucon.get_multiple(params)`: Get values for multiple specified parameters.
-- `Nucon.set_dummy_mode(dummy_mode)`: Enable or disable dummy mode for testing
+- `nucon.get_multiple_iter(params)`: Get values for multiple specified parameters as a generator.
+- `nucon.get_game_variable_names()`: Query the game for all exposed variable names (GET and POST), excluding special endpoints.
+- `nucon.set_dummy_mode(dummy_mode)`: In dummy mode, returns sensible values without connecting to the game and silently ignores writes.
+- `nucon.set_cheat_mode(cheat_mode)`: Enable writing to cheat parameters (`FUN_*` event triggers). Default `False`.
+
+Valve API (motorized actuators: OPEN/CLOSE powers the motor, OFF holds current position):
+- `nucon.get_valve(name)`: Get state dict for a single valve (`Value`, `IsOpened`, `IsClosed`, `Stuck`, …).
+- `nucon.get_valves()`: Get state dict for all 53 valves.
+- `nucon.open_valve(name)` / `nucon.open_valves(names)`: Power actuator toward open.
+- `nucon.close_valve(name)` / `nucon.close_valves(names)`: Power actuator toward closed.
+- `nucon.off_valve(name)` / `nucon.off_valves(names)`: Cut actuator power, hold current position (normal resting state).
 
 Custom Enum Types:
-- `PumpStatus`: Enum for pump status (INACTIVE, ACTIVE_NO_SPEED_REACHED, ACTIVE_SPEED_REACHED, REQUIRES_MAINTENANCE, NOT_INSTALLED, INSUFFICIENT_ENERGY)
+- `PumpStatus`: Enum for pump status (INACTIVE, ACTIVE_NO_SPEED_REACHED\*, ACTIVE_SPEED_REACHED\*, REQUIRES_MAINTENANCE, NOT_INSTALLED, INSUFFICIENT_ENERGY)
-- `PumpDryStatus`: Enum for pump dry status (ACTIVE_WITHOUT_FLUID, INACTIVE_OR_ACTIVE_WITH_FLUID)
+- `PumpDryStatus`: Enum for pump dry status (ACTIVE_WITHOUT_FLUID\*, INACTIVE_OR_ACTIVE_WITH_FLUID)
-- `PumpOverloadStatus`: Enum for pump overload status (ACTIVE_AND_OVERLOAD, INACTIVE_OR_ACTIVE_NO_OVERLOAD)
+- `PumpOverloadStatus`: Enum for pump overload status (ACTIVE_AND_OVERLOAD\*, INACTIVE_OR_ACTIVE_NO_OVERLOAD)
-- `BreakerStatus`: Enum for breaker status (OPEN, CLOSED)
+- `BreakerStatus`: Enum for breaker status (OPEN\*, CLOSED)
 
-## Reinforcement Learning (Work in Progress)
+\*: Truthy value (will be treated as true in e.g. if statements).
 
-NuCon includes a preliminary Reinforcement Learning (RL) environment based on the OpenAI Gym interface. This feature is currently a work in progress and requires additional dependencies.
+So if you're not in the mood to play the game manually, this API can be used to easily create your own automations and control systems. Maybe a little PID controller for the rods — or a full classical reactor operator with grid-demand following, pressurizer control, and a live TUI, like the one in `scripts/reactor_control.py`? Or, if you wanna go crazy, why not try some
+
+## Reinforcement Learning
+
+NuCon includes a Reinforcement Learning (RL) environment based on the OpenAI Gym interface. This allows you to train control policies for the Nucleares game instead of writing them yourself. Requires additional dependencies.
 
 ### Additional Dependencies
 
-To use the RL features, you'll need to install the following packages:
+To use you'll need to install `gymnasium` and `numpy`. You can do so via
 
 ```bash
-pip install gymnasium numpy
+pip install -e '.[rl]'
 ```
 
-### RL Environment
+### Environments
 
-The `NuconEnv` class in `nucon/rl.py` provides a Gym-compatible environment for reinforcement learning tasks in the Nucleares simulation. Key features include:
+Two environment classes are provided in `nucon/rl.py`:
 
-- Observation space: Includes all readable parameters from the Nucon system.
+**`NuconEnv`**: classic fixed-objective environment. You define one or more objectives at construction time (e.g. maximise power output, keep temperature in range). The agent always trains toward the same goal.
-- Action space: Encompasses all writable parameters in the Nucon system.
-- Step function: Applies actions to the Nucon system and returns new observations.
-- Objective function: Allows for custom objective functions to be defined for training.
 
-### Usage
+- Observation space: all readable numeric parameters (~290 dims).
+- Action space: all readable-back writable parameters (~30 dims): 9 individual rod bank positions, 3 MSCVs, 3 turbine bypass valves, 6 coolant pump speeds, condenser pump, freight/vent switches, resistor banks, and more.
+- Objectives: predefined strings (`'max_power'`, `'episode_time'`) or arbitrary callables `(obs) -> float`. Multiple objectives are weighted-summed.
+
+**`NuconGoalEnv`**: goal-conditioned environment. The desired goal (e.g. target generator output) is sampled at the start of each episode and provided as part of the observation. A single policy learns to reach *any* goal in the specified range, making it far more useful than a fixed-objective agent. Designed for training with [Hindsight Experience Replay (HER)](https://arxiv.org/abs/1707.01495), which makes sparse-reward goal-conditioned training tractable.
+
+- Observation space: `Dict` with keys `observation` (non-goal params), `achieved_goal` (current goal param values, normalised to [0,1]), `desired_goal` (target, normalised to [0,1]).
+- Goals are sampled uniformly from the specified `goal_range` each episode.
+- Reward defaults to negative L2 distance in normalised goal space (dense). Pass `tolerance` for a sparse `{0, -1}` reward; this works particularly well with HER.
+
+### NuconEnv Usage
 
-Here's a basic example of how to use the RL environment:
 ```python
-from nucon.rl import NuconEnv
+from nucon.rl import NuconEnv, Parameterized_Objectives
 
 env = NuconEnv(objectives=['max_power'], seconds_per_step=5)
+# env2 = gym.make('Nucon-max_power-v0')
+# env3 = NuconEnv(objectives=[Parameterized_Objectives['target_temperature'](goal_temp=350)], objective_weights=[1.0], seconds_per_step=5)
 
 obs, info = env.reset()
 for _ in range(1000):
@@ -122,17 +155,298 @@ for _ in range(1000):
 env.close()
 ```
 
+Objectives takes either strings of the name of predefined objectives, or lambda functions which take an observation and return a scalar reward. Final rewards are (weighted) summed across all objectives. `info['objectives']` contains all objectives and their values.
+
+You can e.g. train a PPO agent using the [sb3](https://github.com/DLR-RM/stable-baselines3) implementation:
+```python
+from nucon.rl import NuconEnv
+from stable_baselines3 import PPO
+
+env = NuconEnv(objectives=['max_power'], seconds_per_step=5)
+
+model = PPO(
+    "MlpPolicy",
+    env,
+    verbose=1,
+    learning_rate=3e-4,
+    n_steps=2048,
+    batch_size=64,
+    n_epochs=10,
+    gamma=0.99,
+    gae_lambda=0.95,
+    clip_range=0.2,
+    ent_coef=0.01,
+)
+model.learn(total_timesteps=100_000)
+
+obs, info = env.reset()
+for _ in range(1000):
+    action, _states = model.predict(obs, deterministic=True)
+    obs, reward, terminated, truncated, info = env.step(action)
+    if terminated or truncated:
+        obs, info = env.reset()
+env.close()
+```
+
+### NuconGoalEnv + HER Usage
+
+HER works by relabelling past trajectories with the goal that was *actually achieved*, turning every episode into useful training signal even when the agent never reaches the intended target. This makes it much more sample-efficient than standard RL for goal-reaching tasks. This matters a lot given how slow the real game is.
+
+```python
+from nucon.rl import NuconGoalEnv, Parameterized_Objectives, Parameterized_Terminators
+from stable_baselines3 import SAC
+from stable_baselines3.her.her_replay_buffer import HerReplayBuffer
+
+
+env = NuconGoalEnv(
+    goal_params=['GENERATOR_0_KW', 'GENERATOR_1_KW', 'GENERATOR_2_KW'],
+    goal_range={
+        'GENERATOR_0_KW': (0.0, 1200.0),
+        'GENERATOR_1_KW': (0.0, 1200.0),
+        'GENERATOR_2_KW': (0.0, 1200.0),
+    },
+    tolerance=0.05,        # sparse: within 5% of range counts as success (recommended with HER)
+    seconds_per_step=5,
+    simulator=simulator,   # use a pre-trained simulator for fast pre-training
+    # Keep policy within the simulator's known data distribution.
+    # SIM_UNCERTAINTY (kNN-GP posterior std) is injected into obs when a simulator is active.
+    # Tune start/scale/threshold to taste.
+    additional_objectives=[Parameterized_Objectives['uncertainty_penalty'](start=0.3, scale=1.0)],
+    terminators=[Parameterized_Terminators['uncertainty_abort'](threshold=0.7)],
+)
+# Or use a preset: env = gym.make('Nucon-goal_power-v0', simulator=simulator)
+
+model = SAC(
+    'MultiInputPolicy',
+    env,
+    replay_buffer_class=HerReplayBuffer,
+    replay_buffer_kwargs={'n_sampled_goal': 4, 'goal_selection_strategy': 'future'},
+    verbose=1,
+    learning_rate=1e-3,
+    batch_size=256,
+    tau=0.005,
+    gamma=0.98,
+    train_freq=1,
+    gradient_steps=1,
+)
+model.learn(total_timesteps=500_000)
+```
+
+At inference time, inject any target by constructing the observation manually:
+```python
+import numpy as np
+obs, _ = env.reset()
+# Override the desired goal (values are normalised to [0,1] within goal_range)
+obs['desired_goal'] = np.array([0.8, 0.8, 0.8], dtype=np.float32)  # ~960 kW per generator
+action, _ = model.predict(obs, deterministic=True)
+```
+
+Predefined goal environments:
+- `Nucon-goal_power-v0`: target total generator output (3 × 0–1200 kW)
+- `Nucon-goal_temp-v0`: target core temperature (280–380 °C)
+
+RL algorithms require a huge number of training steps, and Nucleares is slow and cannot be trivially parallelised. That's why NuCon provides a built-in simulator.
+
+## Simulator
+
+NuCon provides a built-in simulator to address the challenge of slow training times in the actual Nucleares game. This simulator allows for rapid prototyping and testing of control policies without the need for the full game environment. Key features include:
+
+- Mimics the behavior of the Nucleares game API
+- Configurable initial states and operating modes
+- Faster than real-time simulation
+- Supports parallel execution for increased training throughput
+
+### Additional Dependencies
+
+To use you'll need to install `torch` and `flask`. You can do so via
+
+```bash
+pip install -e '.[sim]'
+```
+
+### Usage
+
+To use the NuCon simulator:
+
+```python
+from nucon import Nucon
+from nucon.sim import NuconSimulator, OperatingState
+
+# Create a simulator instance
+simulator = NuconSimulator()
+
+# Load a dynamics model (explained later)
+simulator.load_model('path/to/model.pth')
+
+# Set initial state (optional)
+simulator.set_state(OperatingState.NOMINAL)
+
+# The web server starts automatically in __init__; access via nucon using the simulator's port
+nucon = Nucon(port=simulator.port)
+
+# Or use the simulator with NuconEnv
+from nucon.rl import NuconEnv
+env = NuconEnv(simulator=simulator) # When given a similator, instead of waiting on the game, we will tell the simulator to skip forward after each step
+
+# Train your RL agent using the simulator
+# ...
+```
+
+The simulator needs an accurate dynamics model of the game. NuCon provides tools to learn one from real gameplay data.
+
+## Model Learning
+
+To address the challenge of unknown game dynamics, NuCon provides tools for collecting data, creating datasets, and training models to learn the reactor dynamics. Key features include:
+
+- **Data Collection**: Gathers state transitions from human play or automated agents. `time_delta` is specified in game-time seconds; wall-clock sleep is automatically adjusted for `GAME_SIM_SPEED` so collected deltas are uniform regardless of simulation speed.
+- **Automatic param filtering**: Junk params (GAME_VERSION, TIME, ALARMS_ACTIVE, …) and params from uninstalled subsystems (returns `None`) are automatically excluded from model inputs/outputs.
+- **Two model backends**: Neural network (NN) or a local Gaussian Process approximated via k-Nearest Neighbours (kNN-GP).
+- **Uncertainty estimation**: The kNN-GP backend returns a GP posterior standard deviation alongside each prediction; 0 means the query lies on known data, ~1 means it is out of distribution.
+- **Dataset management**: Tools for saving, loading, merging, and pruning datasets.
+
+### Additional Dependencies
+
+```bash
+pip install -e '.[model]'
+```
+
+### Model selection
+
+**kNN-GP** (the `ReactorKNNModel` backend) is a local Gaussian Process: it finds the `k` nearest neighbours in the training set, fits an RBF kernel on them, and returns a prediction plus a GP posterior std as uncertainty. It works well from a few hundred samples and requires no training. **NN** needs input normalisation and several thousand samples to generalise; use it once you have a large dataset. For initial experiments, start with kNN-GP (`k=10`).
+
+### Usage
+
+```python
+from nucon.model import NuconModelLearner
+
+# --- Data collection ---
+learner = NuconModelLearner(
+    time_delta=10.0,             # 10 game-seconds per step (wall sleep auto-scales with sim speed)
+    include_valve_states=False,  # set True to include all 53 valve positions as model inputs
+)
+learner.collect_data(num_steps=1000)
+learner.save_dataset('reactor_dataset.pkl')
+
+# Merge datasets collected across multiple sessions
+learner.merge_datasets('other_session.pkl')
+
+# --- Neural network backend ---
+nn_learner = NuconModelLearner(dataset_path='reactor_dataset.pkl')
+nn_learner.train_model(batch_size=32, num_epochs=50)  # creates NN model on first call
+# Drop samples the NN already predicts well (keep hard cases for further training)
+nn_learner.drop_well_fitted(error_threshold=1.0)
+nn_learner.save_model('reactor_nn.pth')
+
+# --- kNN-GP backend ---
+knn_learner = NuconModelLearner(dataset_path='reactor_dataset.pkl')
+# Drop near-duplicate samples before fitting (keeps diverse coverage).
+# A sample is dropped only if BOTH its input state AND output transition
+# are within the given distances of an already-kept sample.
+knn_learner.drop_redundant(min_state_distance=0.1, min_output_distance=0.05)
+knn_learner.fit_knn(k=10)  # creates kNN-GP model on first call
+
+# Point prediction
+state = knn_learner._get_state()
+pred = knn_learner.model.forward(state, time_delta=10.0)
+
+# Prediction with uncertainty
+pred, uncertainty = knn_learner.predict_with_uncertainty(state, time_delta=10.0)
+print(f"CORE_TEMP: {pred['CORE_TEMP']:.1f} ± {uncertainty:.3f} (std, GP posterior)")
+# uncertainty ≈ 0: confident (query near known data)
+# uncertainty ≈ 1: out of distribution
+
+knn_learner.save_model('reactor_knn.pkl')
+```
+
+The trained models can be integrated into the NuconSimulator to provide accurate dynamics based on real game data.
+
+## Full Training Loop
+
+The recommended end-to-end workflow for training an RL operator is an iterative cycle of real-game data collection, model fitting, and simulated training. The real game is slow and cannot be parallelised, so the bulk of RL training happens in the simulator. The game is used only as an oracle for data and evaluation.
+
+```
+┌─────────────────────────────────────────────────────────────┐
+│  1. Human dataset collection                                │
+│     Play the game: start up the reactor, operate it across  │
+│     a range of states. NuCon records state transitions.     │
+└───────────────────────┬─────────────────────────────────────┘
+                        │
+                        ▼
+┌─────────────────────────────────────────────────────────────┐
+│  2. Initial model fitting                                   │
+│     Fit NN or kNN dynamics model to the collected dataset.  │
+│     kNN is instant; NN needs gradient steps but generalises │
+│     better with more data.                                  │
+└───────────────────────┬─────────────────────────────────────┘
+                        │
+              ┌─────────▼──────────┐
+              │   3. Train RL      │◄───────────────────────┐
+              │   in simulator     │                        │
+              │   (fast, many      │                        │
+              │   trajectories)    │                        │
+              └─────────┬──────────┘                        │
+                        │                                   │
+                        ▼                                   │
+              ┌─────────────────────┐                       │
+              │  4. Eval in game    │                       │
+              │  + collect new data │                       │
+              │  (merge & prune     │                       │
+              │   dataset)          │                       │
+              └─────────┬───────────┘                       │
+                        │                                   │
+                        ▼                                   │
+              ┌─────────────────────┐      model improved?  │
+              │  5. Refit model     ├──────── yes ──────────┘
+              │  on expanded data   │
+              └─────────────────────┘
+```
+
+**Step 1 — Human dataset collection**: Run `scripts/collect_dataset.py` during your play session (see [Scripts](#scripts)). Cover a wide range of states: startup from cold, ramping power, individual rod bank adjustments. Diversity in the dataset directly determines simulator accuracy. See [Model Learning](#model-learning) for collection details.
+
+**Step 2 — Initial model fitting**: Fit a kNN-GP model (instant) or NN (better extrapolation with larger datasets) using `fit_knn()` or `train_model()`. Prune near-duplicate samples with `drop_redundant()` before fitting. See [Model Learning](#model-learning).
+
+**Step 3 — Train RL in simulator**: Load the fitted model into `NuconSimulator`, then train a `NuconGoalEnv` policy with SAC + HER. The simulator runs far faster than the real game, allowing many trajectories in reasonable time. Pass `Parameterized_Objectives['uncertainty_penalty']` and `Parameterized_Terminators['uncertainty_abort']` as additional objectives/terminators to discourage the policy from wandering into regions the model hasn't seen; `SIM_UNCERTAINTY` is automatically injected into the obs dict when a simulator is active. See [NuconGoalEnv + HER Usage](#nucongoalenv--her-usage) and `scripts/train_sac.py` for a complete example.
+
+**Step 4 — Eval in game + collect new data**: Run the trained policy against the real game. This validates simulator accuracy and simultaneously collects new data from states the policy visits, which may be regions the original dataset missed. Run a second `NuconModelLearner` in a background thread to collect concurrently.
+
+**Step 5 — Refit model on expanded data**: Merge new data into the original dataset with `merge_datasets()`, prune with `drop_redundant()`, and refit. Then return to Step 3 with the improved model. Each iteration the simulator gets more accurate and the policy improves.
+
+Stop when the policy performs well in the real game and kNN-GP uncertainty stays low throughout an episode, indicating the policy stays within the known data distribution.
+
+## Scripts
+
+Ready-to-run scripts in the `scripts/` directory covering the most common workflows.
+
+**`scripts/collect_dataset.py`** — collect a dynamics dataset while playing the game:
+```bash
+python scripts/collect_dataset.py --steps 1000 --delta 10 --out reactor_dataset.pkl
+# Ctrl-C to stop early; data is saved on exit
+# Merge a previous session: --merge previous.pkl
+```
+
+**`scripts/train_sac.py`** — train a SAC + HER goal-conditioned policy on the kNN-GP simulator:
+```bash
+python scripts/train_sac.py
+# Expects /tmp/reactor_knn.pkl and /tmp/nucon_dataset.pkl
+# Saves trained policy to /tmp/sac_nucon_knn.zip
+```
+This script is the most elaborate end-to-end example: it loads a pre-fitted kNN-GP model, seeds episode resets from dataset states, uses delta actions and an uncertainty penalty, and configures SAC + HER for fast sim training.
+
 ## Testing
 
-NuCon includes a test suite to verify its functionality and compatibility with the Nucleares game. 
+NuCon includes a test suite to verify its functionality and compatibility with the Nucleares game.
 
 ### Running Tests
 
 To run the tests:
 
-1. Ensure the Nucleares game is running and accessible at http://localhost:8080/ (or update the URL in the test setup).
+1. Ensure the Nucleares game is running and accessible at http://localhost:8785/ (or update the URL in the test setup).
-2. Install pytest: `pip install pytest`
+2. Install pytest: `pip install pytest` (or `pip install -e .[dev]`)
-3. Run the tests: `pytest test/test.py`
+3. Run the tests:
+   ```bash
+   pytest test/test_core.py
+   pytest test/test_sim.py
+   ```
 
 ### Test Coverage
 
@@ -142,6 +456,54 @@ The tests verify:
 - Writable parameters can be written to
 - Non-writable parameters cannot be written to, even when force-writing
 - Enum parameters and their custom truthy values behave correctly
+- Simulator functionality and consistency
+
+---
+
+![NuCon Meme](README_meme.jpg)
+
+To use you'll need to install `pillow`. You can do so via
+
+```bash
+pip install -e '.[drake]'
+```
+
+### Usage:
+
+```python
+from nucon.drake import create_drake_meme
+
+items = [
+    (False, "Play Nucleares manually"),
+    (True, "Automate it with a script"),
+    (False, "But the web interface is tedious to use"),
+    (True, "Write an elegant libary to interface with the game and then use that to write the script"),
+    (False, "But I would still need to write the control policy by hand"),
+    (True, "Let's extend the libary such that it trains a policy via Reinforcement Learning"),
+    (False, "But RL takes a huge number of training samples"),
+    (True, "Extend the libary to also include an efficient simulator"),
+    (False, "But I don't know what the actual internal dynamics are"),
+    (True, "Extend the libary once more to also include a neural network dynamics model"),
+    (True, "And I'm gonna put a drake meme on the README"),
+    (False, "Online meme generators only support a single yes/no pair"),
+    (True, "Let's also add a drake meme generator to the libary"),
+    ]
+
+meme = create_drake_meme(items)
+meme.save("README_meme.jpg")
+```
 
 ## Disclaimer
 NuCon is an unofficial tool and is not affiliated with or endorsed by the creators of Nucleares.
+
+## Citing
+What? Why would you wanna cite it? What are you even doing?
+```
+@misc{nucon,
+  title = {NuCon},
+  author = {Dominik Roth},
+  abstract = {NuCon is a Python library to interface with and control Nucleares, a nuclear reactor simulation game. Includes gymnasium bindings for Reinforcement Learning.},
+  url = {https://git.dominik-roth.eu/dodox/NuCon},
+  year = {2024},
+}
+```

logo.svg

@@ -0,0 +1,27 @@
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 512 512" width="512" height="512">
+  <!-- Background -->
+  <rect width="512" height="512" fill="#0d0d0d" rx="72"/>
+
+  <!-- Outer pressure vessel ring -->
+  <circle cx="256" cy="256" r="185" fill="none" stroke="#3ab5f0" stroke-width="10"/>
+
+  <!-- Inner core ring -->
+  <circle cx="256" cy="256" r="65" fill="none" stroke="#3ab5f0" stroke-width="10"/>
+
+  <!-- Central nucleus -->
+  <circle cx="256" cy="256" r="18" fill="#3ab5f0"/>
+
+  <!-- 6 control rods: gap inside core, line from core to vessel -->
+  <!-- θ=0° (top) -->
+  <line x1="256" y1="191" x2="256" y2="71"  stroke="#3ab5f0" stroke-width="9" stroke-linecap="round"/>
+  <!-- θ=60° -->
+  <line x1="312" y1="223" x2="416" y2="163" stroke="#3ab5f0" stroke-width="9" stroke-linecap="round"/>
+  <!-- θ=120° -->
+  <line x1="312" y1="289" x2="416" y2="349" stroke="#3ab5f0" stroke-width="9" stroke-linecap="round"/>
+  <!-- θ=180° (bottom) -->
+  <line x1="256" y1="321" x2="256" y2="441" stroke="#3ab5f0" stroke-width="9" stroke-linecap="round"/>
+  <!-- θ=240° -->
+  <line x1="200" y1="289" x2="96"  y2="349" stroke="#3ab5f0" stroke-width="9" stroke-linecap="round"/>
+  <!-- θ=300° -->
+  <line x1="200" y1="223" x2="96"  y2="163" stroke="#3ab5f0" stroke-width="9" stroke-linecap="round"/>
+</svg>

nucon/__init__.py

@@ -1 +1 @@
-from nucon.core import *
+from nucon.core import Nucon, ParameterEnum, PumpStatus, PumpDryStatus, PumpOverloadStatus, CoreState, CoolantCoreState, RodsState, NuconParameter, BreakerStatus

nucon/core.py

@@ -1,12 +1,23 @@
-from enum import Enum, IntEnum
+from enum import Enum
+from typing import Union, Dict, Type, List, Optional, Any, Iterator, Tuple
 import requests
-from typing import Union, Dict, Type, List
+import random
 
-class NuconConfig:
+class ParameterEnum(Enum):
-    base_url = "http://localhost:8080/"
+    @classmethod
-    dummy_mode = False
+    def _missing_(cls, value: Any) -> Union['ParameterEnum', None]:
+        if isinstance(value, str):
+            if value.lower() == 'true':
+                return cls(True)
+            elif value.lower() == 'false':
+                return cls(False)
+            try:
+                return cls(int(value))
+            except ValueError:
+                pass
+        return None
 
-class PumpStatus(IntEnum):
+class PumpStatus(ParameterEnum):
     INACTIVE = 0
     ACTIVE_NO_SPEED_REACHED = 1
     ACTIVE_SPEED_REACHED = 2
@@ -17,228 +28,444 @@ class PumpStatus(IntEnum):
     def __bool__(self):
         return self.value in (1, 2)
 
-class PumpDryStatus(IntEnum):
+class PumpDryStatus(ParameterEnum):
     ACTIVE_WITHOUT_FLUID = 1
     INACTIVE_OR_ACTIVE_WITH_FLUID = 4
 
     def __bool__(self):
-        return self.value == 4
+        return self.value == 1
 
-class PumpOverloadStatus(IntEnum):
+class PumpOverloadStatus(ParameterEnum):
     ACTIVE_AND_OVERLOAD = 1
     INACTIVE_OR_ACTIVE_NO_OVERLOAD = 4
 
     def __bool__(self):
-        return self.value == 4
+        return self.value == 1
 
-class BreakerStatus(Enum):
+class BreakerStatus(ParameterEnum):
     OPEN = True
     CLOSED = False
 
     def __bool__(self):
         return self.value
 
-class Nucon(Enum):
+CoreState = str
-    CORE_TEMP = ("CORE_TEMP", float, False)
+CoolantCoreState = str
-    CORE_TEMP_OPERATIVE = ("CORE_TEMP_OPERATIVE", float, False)
+RodsState = str
-    CORE_TEMP_MAX = ("CORE_TEMP_MAX", float, False)
-    CORE_TEMP_MIN = ("CORE_TEMP_MIN", float, False)
-    CORE_TEMP_RESIDUAL = ("CORE_TEMP_RESIDUAL", float, False)
-    CORE_PRESSURE = ("CORE_PRESSURE", float, False)
-    CORE_PRESSURE_MAX = ("CORE_PRESSURE_MAX", float, False)
-    CORE_PRESSURE_OPERATIVE = ("CORE_PRESSURE_OPERATIVE", float, False)
-    CORE_INTEGRITY = ("CORE_INTEGRITY", float, False)
-    CORE_WEAR = ("CORE_WEAR", float, False)
-    CORE_STATE = ("CORE_STATE", int, False)
-    CORE_STATE_CRITICALITY = ("CORE_STATE_CRITICALITY", float, False)
-    CORE_CRITICAL_MASS_REACHED = ("CORE_CRITICAL_MASS_REACHED", bool, False)
-    CORE_CRITICAL_MASS_REACHED_COUNTER = ("CORE_CRITICAL_MASS_REACHED_COUNTER", int, False)
-    CORE_IMMINENT_FUSION = ("CORE_IMMINENT_FUSION", bool, False)
-    CORE_READY_FOR_START = ("CORE_READY_FOR_START", bool, False)
-    CORE_STEAM_PRESENT = ("CORE_STEAM_PRESENT", bool, False)
-    CORE_HIGH_STEAM_PRESENT = ("CORE_HIGH_STEAM_PRESENT", bool, False)
 
-    TIME = ("TIME", float, False)
+# Known game endpoints that are not regular parameters (JSON/HTML data blobs, meta-commands)
-    TIME_STAMP = ("TIME_STAMP", str, False)
+SPECIAL_VARIABLES = frozenset({
+    'WEBSERVER_BATCH_GET',
+    'WEBSERVER_LIST_VARIABLES',
+    'WEBSERVER_LIST_VARIABLES_JSON',
+    'WEBSERVER_VIEW_VARIABLES',
+    'VALVE_PANEL_JSON',
+    'RESISTOR_BANKS_JSON',
+    'INSTALLED_LOOPS_JSON',
+    'INVENTORY_HTML',
+    'MAINTENANCE_REPORT_HTML',
+    'WEATHER_FORECAST_JSON',
+})
 
-    COOLANT_CORE_STATE = ("COOLANT_CORE_STATE", int, False)
+class NuconParameter:
-    COOLANT_CORE_PRESSURE = ("COOLANT_CORE_PRESSURE", float, False)
+    def __init__(self, nucon: 'Nucon', id: str, param_type: Type, is_writable: bool, min_val: Optional[Union[int, float]] = None, max_val: Optional[Union[int, float]] = None, unit: Optional[str] = None, is_readable: bool = True, is_cheat: bool = False):
-    COOLANT_CORE_MAX_PRESSURE = ("COOLANT_CORE_MAX_PRESSURE", float, False)
+        self.nucon = nucon
-    COOLANT_CORE_VESSEL_TEMPERATURE = ("COOLANT_CORE_VESSEL_TEMPERATURE", float, False)
-    COOLANT_CORE_QUANTITY_IN_VESSEL = ("COOLANT_CORE_QUANTITY_IN_VESSEL", float, False)
-    COOLANT_CORE_PRIMARY_LOOP_LEVEL = ("COOLANT_CORE_PRIMARY_LOOP_LEVEL", float, False)
-    COOLANT_CORE_FLOW_SPEED = ("COOLANT_CORE_FLOW_SPEED", float, False)
-    COOLANT_CORE_FLOW_ORDERED_SPEED = ("COOLANT_CORE_FLOW_ORDERED_SPEED", float, True)
-    COOLANT_CORE_FLOW_REACHED_SPEED = ("COOLANT_CORE_FLOW_REACHED_SPEED", float, False)
-    COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT = ("COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT", int, False)
-    COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT = ("COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT", int, False)
-
-    COOLANT_CORE_CIRCULATION_PUMP_0_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_0_STATUS", PumpStatus, False)
-    COOLANT_CORE_CIRCULATION_PUMP_1_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_1_STATUS", PumpStatus, False)
-    COOLANT_CORE_CIRCULATION_PUMP_2_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_2_STATUS", PumpStatus, False)
-
-    COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS", PumpDryStatus, False)
-    COOLANT_CORE_CIRCULATION_PUMP_1_DRY_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_1_DRY_STATUS", PumpDryStatus, False)
-    COOLANT_CORE_CIRCULATION_PUMP_2_DRY_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_2_DRY_STATUS", PumpDryStatus, False)
-
-    COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS", PumpOverloadStatus, False)
-    COOLANT_CORE_CIRCULATION_PUMP_1_OVERLOAD_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_1_OVERLOAD_STATUS", PumpOverloadStatus, False)
-    COOLANT_CORE_CIRCULATION_PUMP_2_OVERLOAD_STATUS = ("COOLANT_CORE_CIRCULATION_PUMP_2_OVERLOAD_STATUS", PumpOverloadStatus, False)
-
-    COOLANT_CORE_CIRCULATION_PUMP_0_ORDERED_SPEED = ("COOLANT_CORE_CIRCULATION_PUMP_0_ORDERED_SPEED", float, True)
-    COOLANT_CORE_CIRCULATION_PUMP_1_ORDERED_SPEED = ("COOLANT_CORE_CIRCULATION_PUMP_1_ORDERED_SPEED", float, True)
-    COOLANT_CORE_CIRCULATION_PUMP_2_ORDERED_SPEED = ("COOLANT_CORE_CIRCULATION_PUMP_2_ORDERED_SPEED", float, True)
-    COOLANT_CORE_CIRCULATION_PUMP_0_SPEED = ("COOLANT_CORE_CIRCULATION_PUMP_0_SPEED", float, False)
-    COOLANT_CORE_CIRCULATION_PUMP_1_SPEED = ("COOLANT_CORE_CIRCULATION_PUMP_1_SPEED", float, False)
-    COOLANT_CORE_CIRCULATION_PUMP_2_SPEED = ("COOLANT_CORE_CIRCULATION_PUMP_2_SPEED", float, False)
-
-    RODS_STATUS = ("RODS_STATUS", int, False)
-    RODS_MOVEMENT_SPEED = ("RODS_MOVEMENT_SPEED", float, False)
-    RODS_MOVEMENT_SPEED_DECREASED_HIGH_TEMPERATURE = ("RODS_MOVEMENT_SPEED_DECREASED_HIGH_TEMPERATURE", bool, False)
-    RODS_DEFORMED = ("RODS_DEFORMED", bool, False)
-    RODS_TEMPERATURE = ("RODS_TEMPERATURE", float, False)
-    RODS_MAX_TEMPERATURE = ("RODS_MAX_TEMPERATURE", float, False)
-    RODS_POS_ORDERED = ("RODS_POS_ORDERED", float, True)
-    RODS_POS_ACTUAL = ("RODS_POS_ACTUAL", float, False)
-    RODS_POS_REACHED = ("RODS_POS_REACHED", bool, False)
-    RODS_QUANTITY = ("RODS_QUANTITY", int, False)
-    RODS_ALIGNED = ("RODS_ALIGNED", bool, False)
-
-    GENERATOR_0_KW = ("GENERATOR_0_KW", float, False)
-    GENERATOR_1_KW = ("GENERATOR_1_KW", float, False)
-    GENERATOR_2_KW = ("GENERATOR_2_KW", float, False)
-    GENERATOR_0_V = ("GENERATOR_0_V", float, False)
-    GENERATOR_1_V = ("GENERATOR_1_V", float, False)
-    GENERATOR_2_V = ("GENERATOR_2_V", float, False)
-    GENERATOR_0_A = ("GENERATOR_0_A", float, False)
-    GENERATOR_1_A = ("GENERATOR_1_A", float, False)
-    GENERATOR_2_A = ("GENERATOR_2_A", float, False)
-    GENERATOR_0_HERTZ = ("GENERATOR_0_HERTZ", float, False)
-    GENERATOR_1_HERTZ = ("GENERATOR_1_HERTZ", float, False)
-    GENERATOR_2_HERTZ = ("GENERATOR_2_HERTZ", float, False)
-    GENERATOR_0_BREAKER = ("GENERATOR_0_BREAKER", BreakerStatus, True)
-    GENERATOR_1_BREAKER = ("GENERATOR_1_BREAKER", BreakerStatus, True)
-    GENERATOR_2_BREAKER = ("GENERATOR_2_BREAKER", BreakerStatus, True)
-
-    STEAM_TURBINE_0_RPM = ("STEAM_TURBINE_0_RPM", float, False)
-    STEAM_TURBINE_1_RPM = ("STEAM_TURBINE_1_RPM", float, False)
-    STEAM_TURBINE_2_RPM = ("STEAM_TURBINE_2_RPM", float, False)
-    STEAM_TURBINE_0_TEMPERATURE = ("STEAM_TURBINE_0_TEMPERATURE", float, False)
-    STEAM_TURBINE_1_TEMPERATURE = ("STEAM_TURBINE_1_TEMPERATURE", float, False)
-    STEAM_TURBINE_2_TEMPERATURE = ("STEAM_TURBINE_2_TEMPERATURE", float, False)
-    STEAM_TURBINE_0_PRESSURE = ("STEAM_TURBINE_0_PRESSURE", float, False)
-    STEAM_TURBINE_1_PRESSURE = ("STEAM_TURBINE_1_PRESSURE", float, False)
-    STEAM_TURBINE_2_PRESSURE = ("STEAM_TURBINE_2_PRESSURE", float, False)
-
-    def __init__(self, id: str, param_type: Type, is_writable: bool, min_val: Optional[Union[int, float]] = None, max_val: Optional[Union[int, float]] = None):
         self.id = id
         self.param_type = param_type
         self.is_writable = is_writable
+        self.is_readable = is_readable
+        self.is_cheat = is_cheat
         self.min_val = min_val
         self.max_val = max_val
-
+        self.unit = unit
-    def __str__(self):
-        return self.id
 
     @property
     def enum_type(self) -> Type[Enum]:
         return self.param_type if issubclass(self.param_type, Enum) else None
 
-    @property
+    def check_in_range(self, value: Union[int, float, Enum], raise_on_oob: bool = False) -> bool:
-    def value(self) -> Union[float, int, bool, str, Enum]:
+        if self.enum_type:
-        return self.read()
+            if not isinstance(value, self.enum_type):
-
+                if raise_on_oob:
-    def read(self) -> Union[float, int, bool, str, Enum]:
+                    raise ValueError(f"Value {value} is not a valid {self.enum_type.__name__}")
-        return Nucon.get(self)
+                return False
-
+            return True
-    @value.setter
+        
-    def value(self, new_value: Union[float, int, bool, str, Enum]) -> None:
-        self.write(new_value)
-
-    def write(self, new_value: Union[float, int, bool, str, Enum], force: bool = False) -> None:
-        Nucon.set(self, new_value, force)
-
-    def check_in_range(self, value: Union[int, float], raise_on_oob=False) -> None:
         if self.min_val is not None and value < self.min_val:
             if raise_on_oob:
-                raise ValueError(f"Value {value} is below the minimum allowed value {self.min_val} for {self.name}")
+                raise ValueError(f"Value {value} is below the minimum allowed value {self.min_val}")
             return False
         if self.max_val is not None and value > self.max_val:
             if raise_on_oob:
-                raise ValueError(f"Value {value} is above the maximum allowed value {self.max_val} for {self.name}")
+                raise ValueError(f"Value {value} is above the maximum allowed value {self.max_val}")
             return False
         return True
 
-    @classmethod
+    @property
-    def set_base_url(cls, url: str) -> None:
+    def value(self):
-        NuconConfig.base_url = url
+        if not self.is_readable:
+            raise AttributeError(f"Parameter {self.id} is write-only and cannot be read")
+        return self.nucon.get(self)
 
-    @classmethod
+    @value.setter
-    def set_dummy_mode(cls, dummy_mode: bool) -> None:
+    def value(self, new_value):
-        NuconConfig.dummy_mode = dummy_mode
+        self.nucon.set(self, new_value)
 
-    @classmethod
+    def read(self):
-    def get_all_readable(cls) -> List['Nucon']:
+        return self.value
-        return list(cls)  # All parameters are readable
 
-    @classmethod
+    def write(self, new_value, force=False):
-    def get_all_writable(cls) -> List['Nucon']:
+        self.nucon.set(self, new_value, force)
-        return [param for param in cls if param.is_writable]
 
-    @classmethod
+    def __repr__(self):
-    def get(cls, parameter: Union['Nucon', str]) -> Union[float, int, bool, str, Enum]:
+        unit_str = f", unit='{self.unit}'" if self.unit else ""
-        if isinstance(parameter, str):
+        value_str = f", value={self.value}" if self.is_readable else ""
-            parameter = cls[parameter]
+        rw_str = "write-only" if not self.is_readable else f"is_writable={self.is_writable}"
+        admin_str = ", is_cheat=True" if self.is_cheat else ""
+        return f"NuconParameter(id='{self.id}'{value_str}, param_type={self.param_type.__name__}, {rw_str}{admin_str}{unit_str})"
 
-        if NuconConfig.dummy_mode:
+    def __str__(self):
-            return cls._get_dummy_value(parameter)
+        return self.id
 
-        response = requests.get(NuconConfig.base_url, params={"variable": parameter.name})
+class Nucon:
-        
+    def __init__(self, host: str = 'localhost', port: int = 8785, cheat_mode: bool = False):
-        if response.status_code != 200:
+        self.base_url = f'http://{host}:{port}/'
-            raise Exception(f"Failed to query parameter {parameter.name}. Status code: {response.status_code}")
+        self.dummy_mode = False
+        self.cheat_mode = cheat_mode
+        self._parameters = self._create_parameters()
 
-        value = response.text.strip()
+    def _create_parameters(self) -> Dict[str, NuconParameter]:
+        param_values = {
+            # --- Game metadata ---
+            'GAME_VERSION': (str, False),
+            'GAME_DIFFICULTY': (int, False),
+            'TIME': (str, False),
+            'TIME_STAMP': (str, False),
+            'TIME_DAY': (int, False),
+            'ALARMS_ACTIVE': (str, False),
+            'GAME_SIM_SPEED': (float, False),
+            'AMBIENT_TEMPERATURE': (float, False, None, None, '°C'),
+            'FUN_IS_ENABLED': (bool, False),
 
+            # --- Core thermal/pressure ---
+            'CORE_TEMP': (float, False, 0, 1000, '°C'),
+            'CORE_TEMP_OPERATIVE': (float, False, None, None, '°C'),
+            'CORE_TEMP_MAX': (float, False, None, None, '°C'),
+            'CORE_TEMP_MIN': (float, False, None, None, '°C'),
+            'CORE_TEMP_RESIDUAL': (bool, False),
+            'CORE_PRESSURE': (float, False, None, None, 'bar'),
+            'CORE_PRESSURE_MAX': (float, False, None, None, 'bar'),
+            'CORE_PRESSURE_OPERATIVE': (float, False, None, None, 'bar'),
+            'CORE_INTEGRITY': (float, False, 0, 100, '%'),
+            'CORE_WEAR': (float, False, 0, 100, '%'),
+            'CORE_STATE': (CoreState, False),
+            'CORE_STATE_CRITICALITY': (float, False, 0, 1),
+            'CORE_CRITICAL_MASS_REACHED': (bool, False),
+            'CORE_CRITICAL_MASS_REACHED_COUNTER': (int, False),
+            'CORE_IMMINENT_FUSION': (bool, False),
+            'CORE_READY_FOR_START': (bool, False),
+            'CORE_STEAM_PRESENT': (bool, False),
+            'CORE_HIGH_STEAM_PRESENT': (bool, False),
+
+            # --- Core physics ---
+            'CORE_FACTOR': (float, False),
+            'CORE_FACTOR_CHANGE': (float, False),
+            'CORE_OPERATION_MODE': (str, True),
+            'CORE_IODINE_GENERATION': (float, False),
+            'CORE_IODINE_CUMULATIVE': (float, False),
+            'CORE_XENON_GENERATION': (float, False),
+            'CORE_XENON_CUMULATIVE': (float, False),
+
+            # --- Core fuel / bays (9 bays) ---
+            'CORE_FUEL_AVG_FISSIONABLE': (float, False, 0, 100, '%'),
+            'CORE_FUEL_AVG_TEMPERATURE': (float, False, None, None, '°C'),
+            'CORE_FUEL_AVG_POWER_FACTOR': (float, False, 0, 1),
+            **{f'CORE_FUEL_{i}_TEMPERATURE': (float, False, None, None, '°C') for i in range(1, 10)},
+            **{f'CORE_FUEL_{i}_FISSIONABLE': (float, False, 0, 100, '%') for i in range(1, 10)},
+            **{f'CORE_FUEL_{i}_POWER_FACTOR': (float, False, 0, 1) for i in range(1, 10)},
+            **{f'CORE_BAY_{i}_STATE': (str, False) for i in range(1, 10)},
+            **{f'CORE_BAY_{i}_HATCH_OPEN': (bool, False) for i in range(1, 10)},
+
+            # --- Core pool ---
+            'CORE_POOL_PUMP': (int, False),  # command-style (LOAD/OFF/REMOVE), not freely settable
+            'CORE_POOL_COOLANT_TANK_VOLUME': (float, False),
+            'CORE_PRIMARY_CIRCUIT_COOLING_TANK_VOLUME': (float, False),
+            'CORE_EXTERNAL_COOLANT_RESERVOIR_VOLUME': (float, False),
+
+            # --- Primary coolant (core loop) ---
+            'COOLANT_CORE_STATE': (CoolantCoreState, False),
+            'COOLANT_CORE_PRESSURE': (float, False, None, None, 'bar'),
+            'COOLANT_CORE_MAX_PRESSURE': (float, False, None, None, 'bar'),
+            'COOLANT_CORE_VESSEL_TEMPERATURE': (float, False, None, None, '°C'),
+            'COOLANT_CORE_QUANTITY_IN_VESSEL': (float, False),
+            'COOLANT_CORE_PRIMARY_LOOP_LEVEL': (float, False, 0, 100, '%'),
+            'COOLANT_CORE_FLOW_IN': (float, False),
+            'COOLANT_CORE_FLOW_OUT': (float, False),
+            'COOLANT_CORE_FLOW_SPEED': (float, False),
+            'COOLANT_CORE_FLOW_ORDERED_SPEED': (float, False),
+            'COOLANT_CORE_FLOW_REACHED_SPEED': (bool, False),
+            'COOLANT_CORE_QUANTITY_CIRCULATION_PUMPS_PRESENT': (int, False),
+            'COOLANT_CORE_QUANTITY_FREIGHT_PUMPS_PRESENT': (int, False),
+            **{f'COOLANT_CORE_CIRCULATION_PUMP_{i}_STATUS': (PumpStatus, False) for i in range(3)},
+            **{f'COOLANT_CORE_CIRCULATION_PUMP_{i}_DRY_STATUS': (PumpDryStatus, False) for i in range(3)},
+            **{f'COOLANT_CORE_CIRCULATION_PUMP_{i}_OVERLOAD_STATUS': (PumpOverloadStatus, False) for i in range(3)},
+            **{f'COOLANT_CORE_CIRCULATION_PUMP_{i}_ORDERED_SPEED': (float, True, 0, 100) for i in range(3)},
+            **{f'COOLANT_CORE_CIRCULATION_PUMP_{i}_SPEED': (float, False, 0, 100) for i in range(3)},
+            **{f'COOLANT_CORE_CIRCULATION_PUMP_{i}_CAPACITY': (float, False) for i in range(3)},
+
+            # --- Secondary coolant loops ---
+            **{f'COOLANT_SEC_CIRCULATION_PUMP_{i}_STATUS': (PumpStatus, False) for i in range(3)},
+            **{f'COOLANT_SEC_CIRCULATION_PUMP_{i}_DRY_STATUS': (PumpDryStatus, False) for i in range(3)},
+            **{f'COOLANT_SEC_CIRCULATION_PUMP_{i}_OVERLOAD_STATUS': (PumpOverloadStatus, False) for i in range(3)},
+            **{f'COOLANT_SEC_CIRCULATION_PUMP_{i}_ORDERED_SPEED': (float, True, 0, 100) for i in range(3)},
+            **{f'COOLANT_SEC_CIRCULATION_PUMP_{i}_SPEED': (float, False, 0, 100) for i in range(3)},
+            **{f'COOLANT_SEC_CIRCULATION_PUMP_{i}_CAPACITY': (float, False) for i in range(3)},
+            **{f'COOLANT_SEC_{i}_VOLUME': (float, False) for i in range(3)},
+            **{f'COOLANT_SEC_{i}_LIQUID_VOLUME': (float, False) for i in range(3)},
+            **{f'COOLANT_SEC_{i}_PRESSURE': (float, False, None, None, 'bar') for i in range(3)},
+            **{f'COOLANT_SEC_{i}_TEMPERATURE': (float, False, None, None, '°C') for i in range(3)},
+
+            # --- Control rods ---
+            'RODS_STATUS': (RodsState, False),
+            'RODS_MOVEMENT_SPEED': (float, False),
+            'RODS_MOVEMENT_SPEED_DECREASED_HIGH_TEMPERATURE': (bool, False),
+            'RODS_DEFORMED': (bool, False),
+            'RODS_TEMPERATURE': (float, False, None, None, '°C'),
+            'RODS_MAX_TEMPERATURE': (float, False, None, None, '°C'),
+            'RODS_POS_ORDERED': (float, False, 0, 100, '%'),
+            'RODS_POS_ACTUAL': (float, False, 0, 100, '%'),
+            'RODS_POS_REACHED': (bool, False),
+            'RODS_QUANTITY': (int, False),
+            'RODS_ALIGNED': (int, False),
+            **{f'ROD_BANK_POS_{i}_ORDERED': (float, True, 0, 100, '%') for i in range(9)},
+            **{f'ROD_BANK_POS_{i}_ACTUAL': (float, False, 0, 100, '%') for i in range(9)},
+
+            # --- Steam generators ---
+            **{f'STEAM_GEN_{i}_STATUS': (PumpStatus, False) for i in range(3)},
+            **{f'STEAM_GEN_{i}_OUTLET': (float, False) for i in range(3)},
+            **{f'STEAM_GEN_{i}_EVAPORATED': (float, False) for i in range(3)},
+            **{f'STEAM_GEN_{i}_BOILING_POINT': (float, False, None, None, '°C') for i in range(3)},
+            **{f'STEAM_GEN_{i}_INLET': (float, False) for i in range(3)},
+            **{f'STEAM_GEN_{i}_RETURN_FLOW_PLUS_CONDENSED': (float, False) for i in range(3)},
+            **{f'STEAM_GEN_{i}_VENT_SWITCH': (bool, True) for i in range(3)},
+
+            # --- Steam turbines ---
+            **{f'STEAM_TURBINE_{i}_RPM': (float, False) for i in range(3)},
+            **{f'STEAM_TURBINE_{i}_TEMPERATURE': (float, False, None, None, '°C') for i in range(3)},
+            **{f'STEAM_TURBINE_{i}_PRESSURE': (float, False, None, None, 'bar') for i in range(3)},
+            **{f'STEAM_TURBINE_{i}_TORQUE': (float, False) for i in range(3)},
+            **{f'STEAM_TURBINE_{i}_INSTALLED': (bool, False) for i in range(3)},
+            **{f'STEAM_TURBINE_{i}_BYPASS_ACTUAL': (float, False, 0, 100, '%') for i in range(3)},
+
+            # --- MSCV valves ---
+            **{f'MSCV_{i}_OPENING_ACTUAL': (float, False, 0, 100, '%') for i in range(3)},
+
+            # --- Main steam valves ---
+            **{f'VALVE_M0{i}_OPEN': (float, False, 0, 100, '%') for i in range(1, 4)},
+
+            # --- Condenser ---
+            'CONDENSER_VACUUM': (float, False, 0, 100, '%'),
+            'CONDENSER_VACUUM_RELIEF_VALVE_OPENING': (float, False, 0, 100, '%'),
+            'CONDENSER_VACUUM_PUMP_ACTIVE': (bool, False),
+            'CONDENSER_VACUUM_PUMP_MODE': (str, True),
+            'CONDENSER_VACUUM_PUMP_POWER': (float, False, 0, 100, '%'),
+            'CONDENSER_TEMPERATURE': (float, False, None, None, '°C'),
+            'CONDENSER_VOLUME': (float, False),
+            'CONDENSER_VAPOR_VOLUME': (float, False),
+            'CONDENSER_CONDENSATE_FLOW_RATE': (float, False),
+            'CONDENSER_EXTRACTION_FLOW_RATE': (float, False),
+            'CONDENSER_COOLANT_EVAPORATED': (float, False),
+            'CONDENSER_PRESSURE': (float, False, None, None, 'bar'),
+            'CONDENSER_CIRCULATION_PUMP_ACTIVE': (bool, False),
+            'CONDENSER_CIRCULATION_PUMP_OVERLOAD_STATUS': (bool, False),
+            'CONDENSER_CIRCULATION_PUMP_SPEED': (float, False, 0, 100),
+            'CONDENSER_CIRCULATION_PUMP_ORDERED_SPEED': (float, True, 0, 100),
+            'CONDENSER_CIRCULATION_PUMP_SWITCH': (bool, True),
+
+            # --- Vacuum retention / steam ejector ---
+            'VACUUM_RETENTION_TANK_VOLUME': (float, False),
+            'VACUUM_RETENTION_TANK_PRESSURE': (float, False, None, None, 'bar'),
+            'STEAM_EJECTOR_MOTIVE': (float, False),
+            'STEAM_EJECTOR_STARTUP_MOTIVE_VALVE_ORDERED': (float, False, 0, 100, '%'),
+            'STEAM_EJECTOR_STARTUP_MOTIVE_VALVE_ACTUAL': (float, False, 0, 100, '%'),
+            'STEAM_EJECTOR_OPERATIONAL_MOTIVE_VALVE_ORDERED': (float, False, 0, 100, '%'),
+            'STEAM_EJECTOR_OPERATIONAL_MOTIVE_VALVE_ACTUAL': (float, False, 0, 100, '%'),
+            'STEAM_EJECTOR_CONDENSER_RETURN_VALVE_ORDERED': (float, False, 0, 100, '%'),
+            'STEAM_EJECTOR_CONDENSER_RETURN_VALVE_ACTUAL': (float, False, 0, 100, '%'),
+
+            # --- Freight pumps ---
+            'FREIGHT_PUMP_CONDENSER_ACTIVE': (bool, False),
+            'FREIGHT_PUMP_INTERNAL_ACTIVE': (bool, False),
+            'FREIGHT_PUMP_EXTERNAL_ACTIVE': (bool, False),
+            'FREIGHT_PUMP_FEEDWATER_ACTIVE': (bool, False),
+            'FREIGHT_PUMP_CONDENSER_SWITCH': (bool, True),
+            'FREIGHT_PUMP_INTERNAL_SWITCH': (bool, True),
+            'FREIGHT_PUMP_EXTERNAL_SWITCH': (bool, True),
+            'FREIGHT_PUMP_FEEDWATER_SWITCH': (bool, True),
+
+            # --- Generators ---
+            **{f'GENERATOR_{i}_KW': (float, False, None, None, 'kW') for i in range(3)},
+            **{f'GENERATOR_{i}_V': (float, False, None, None, 'V') for i in range(3)},
+            **{f'GENERATOR_{i}_A': (float, False, None, None, 'A') for i in range(3)},
+            **{f'GENERATOR_{i}_HERTZ': (float, False, None, None, 'Hz') for i in range(3)},
+            **{f'GENERATOR_{i}_BREAKER': (BreakerStatus, False) for i in range(3)},
+
+            # --- Resistor banks ---
+            'RES_DIVERT_SURPLUS_FROM_MW': (float, False, None, None, 'MW'),
+            'RES_EFFECTIVELY_DERIVED_ENERGY_MW': (float, False, None, None, 'MW'),
+            'RES_ABSORPTION_CAPACITY_MW': (float, False, None, None, 'MW'),
+            **{f'RESISTOR_BANK_0{i}_SWITCH': (bool, True) for i in range(1, 5)},
+            'RESISTOR_BANKS_MAIN_SWITCH': (bool, True),
+
+            # --- Emergency generators ---
+            'EMERGENCY_GENERATOR_1_MODE': (str, True),
+            'EMERGENCY_GENERATOR_1_STATUS': (str, False),
+            'EMERGENCY_GENERATOR_1_PRESSURIZER': (str, False),
+            'EMERGENCY_GENERATOR_1_FUEL': (float, False),
+            'EMERGENCY_GENERATOR_1_MAINTENANCE_NEEDED': (bool, False),
+            'EMERGENCY_GENERATOR_2_MODE': (str, True),
+            'EMERGENCY_GENERATOR_2_STATUS': (str, False),
+            'EMERGENCY_GENERATOR_2_PRESSURIZER': (str, False),
+            'EMERGENCY_GENERATOR_2_FUEL': (float, False),
+            'EMERGENCY_GENERATOR_2_MAINTENANCE_NEEDED': (bool, False),
+
+            # --- Power ---
+            'POWER_FROM_TURBINE_KW': (float, False, None, None, 'kW'),
+            'POWER_FROM_EXTERNAL_KW': (float, False, None, None, 'kW'),
+            'EMERGENCY_GENERATOR_POWER_OUTPUT_KW': (float, False, None, None, 'kW'),
+            'EMERGENCY_BATTERIES_POWER_OUTPUT_KW': (float, False, None, None, 'kW'),
+            'EMERGENCY_BATTERIES_MODE': (int, True),
+            'POWER_DEMAND_MW': (float, False, None, None, 'MW'),
+            'POWER_MAX_THEORETICAL_FINAL_PLANT_OUTPUT_MW': (float, False, None, None, 'MW'),
+            'POWER_MAX_THEORETICAL_PLANT_OUTPUT_MW': (float, False, None, None, 'MW'),
+
+            # --- Chemistry ---
+            'CHEM_TRUCK_IN_ZONE': (bool, False),
+            'CHEM_TRUCK_CONNECTED': (bool, False),
+            'CHEM_BORON_DOSAGE_ORDERED': (float, False),
+            'CHEM_BORON_DOSAGE_ACTUAL': (float, False),
+            'CHEM_BORON_FILTER_ORDERED': (float, False),
+            'CHEM_BORON_FILTER_ACTUAL': (float, False),
+            'CHEM_BORON_PPM': (float, False, None, None, 'ppm'),
+            'CHEMICAL_DOSING_PUMP_STATUS': (PumpStatus, False),
+            'CHEMICAL_DOSING_PUMP_DRY_STATUS': (PumpDryStatus, False),
+            'CHEMICAL_DOSING_PUMP_OVERLOAD_STATUS': (PumpOverloadStatus, False),
+            'CHEMICAL_FILTER_PUMP_STATUS': (PumpStatus, False),
+            'CHEMICAL_FILTER_PUMP_DRY_STATUS': (PumpDryStatus, False),
+            'CHEMICAL_FILTER_PUMP_OVERLOAD_STATUS': (PumpOverloadStatus, False),
+            'CHEMICAL_CLEANING_PUMP_STATUS': (PumpStatus, False),
+            'CHEMICAL_CLEANING_PUMP_DRY_STATUS': (PumpDryStatus, False),
+            'CHEMICAL_CLEANING_PUMP_OVERLOAD_STATUS': (PumpOverloadStatus, False),
+        }
+
+        # Write-only params: normal operational commands
+        write_only_values = {
+            # --- MSCVs (Main Steam Control Valves) setpoints ---
+            **{f'MSCV_{i}_OPENING_ORDERED': (float, True, 0, 100, '%') for i in range(3)},
+
+            # --- Steam turbine bypass setpoints ---
+            **{f'STEAM_TURBINE_{i}_BYPASS_ORDERED': (float, True, 0, 100, '%') for i in range(3)},
+
+            # --- Steam ejector valve setpoints (0-100 position, not bool) ---
+            'STEAM_EJECTOR_STARTUP_MOTIVE_VALVE': (int, True, 0, 100, '%'),
+            'STEAM_EJECTOR_OPERATIONAL_MOTIVE_VALVE': (int, True, 0, 100, '%'),
+            'STEAM_EJECTOR_CONDENSER_RETURN_VALVE': (int, True, 0, 100, '%'),
+
+            # --- Generic valve commands (value = valve name e.g. "M01", "M02", "M03") ---
+            'VALVE_OPEN': (str, True),
+            'VALVE_CLOSE': (str, True),
+            'VALVE_OFF': (str, True),
+
+            # --- Pump / generator start/stop ---
+            'CONDENSER_VACUUM_PUMP_START_STOP': (bool, True),
+            'EMERGENCY_GENERATOR_1_START_STOP': (bool, True),
+            'EMERGENCY_GENERATOR_2_START_STOP': (bool, True),
+
+            # --- Chemistry setpoints ---
+            'CHEM_BORON_DOSAGE_ORDERED_RATE': (float, True, 0, 100, '%'),
+            'CHEM_BORON_FILTER_ORDERED_SPEED': (float, True, 0, 100, '%'),
+
+            # --- Core safety / operational actions ---
+            'CORE_SCRAM_BUTTON': (bool, True),
+            'CORE_EMERGENCY_STOP': (bool, True),
+            'CORE_END_EMERGENCY_STOP': (bool, True),
+            'RESET_AO': (bool, True),
+            'STEAM_TURBINE_TRIP': (bool, True),
+            'RODS_ALL_POS_ORDERED': (float, True, 0, 100, '%'),
+
+            # --- Core bay physical operations ---
+            **{f'CORE_BAY_{i}_HATCH': (bool, True) for i in range(1, 10)},
+            **{f'CORE_BAY_{i}_FUEL_LOADING': (int, True) for i in range(1, 10)},
+        }
+
+        # Write-only cheat params: game event triggers, blocked unless cheat_mode=True
+        write_only_cheat_values = {
+            'FUN_REQUEST_ENABLE': (bool, True),
+            'FUN_AO_SABOTAGE_ONCE': (bool, True),
+            'FUN_AO_SABOTAGE_TIME': (float, True),
+            'FUN_BANK_ROBBERY': (bool, True),
+            'FUN_BREAKER_TRIP': (bool, True),
+            'FUN_DECREASE_INTEGRITY': (float, True),
+            'FUN_FIRE_DRILL': (bool, True),
+            'FUN_IODINE_SPILL': (bool, True),
+            'FUN_OIL_SPILL': (bool, True),
+            'FUN_PUMP_JAM': (bool, True),
+            'FUN_SHOW_MESSAGE': (str, True),
+            'FUN_TOGGLE_RANDOM_SWITCH': (bool, True),
+            'FUN_TRIGGER_AUDIT': (bool, True),
+            'FUN_WEATHER_CONTROL': (str, True),
+            'FUN_XENON_SPILL': (bool, True),
+        }
+
+        params = {
+            name: NuconParameter(self, name, *values)
+            for name, values in param_values.items()
+        }
+        for name, values in write_only_values.items():
+            params[name] = NuconParameter(self, name, *values, is_readable=False)
+        for name, values in write_only_cheat_values.items():
+            params[name] = NuconParameter(self, name, *values, is_readable=False, is_cheat=True)
+        return params
+
+    def _parse_value(self, parameter: NuconParameter, value: str) -> Union[float, int, bool, str, Enum, None]:
+        if value == '' or value is None or (isinstance(value, str) and value.lower() == 'null'):
+            return None
         if parameter.enum_type:
-            return parameter.enum_type(int(value))
+            try:
-        elif parameter.param_type in (float, int):
+                return parameter.enum_type(value)
-            return parameter.param_type(value)
+            except ValueError as e:
+                raise ValueError(f"Failed to convert {value} to {parameter.enum_type.__name__} for parameter {parameter.id}: {e}")
         elif parameter.param_type == bool:
-            return value.lower() == "true"
+            if isinstance(value, str):
+                if value.lower() not in ('true', 'false'):
+                    raise ValueError(f"Invalid boolean value for parameter {parameter.id}: {value}")
+                return value.lower() == 'true'
+            else:
+                raise ValueError(f"Expected string for boolean parameter {parameter.id}, got {type(value)}")
         else:
-            return value
+            try:
+                return parameter.param_type(value)
+            except ValueError as e:
+                raise ValueError(f"Failed to convert {value} to {parameter.param_type.__name__} for parameter {parameter.id}: {e}")
 
-    @classmethod
+    def get(self, parameter: Union[str, NuconParameter]) -> Union[float, int, bool, str, Enum]:
-    def _get_dummy_value(cls, parameter: 'Nucon') -> Union[float, int, bool, str, Enum]:
-        if parameter.enum_type:
-            return next(iter(parameter.enum_type))
-        elif parameter.param_type == float:
-            return 0.0
-        elif parameter.param_type == int:
-            return 0
-        elif parameter.param_type == bool:
-            return False
-        else:
-            return ""
-
-    @classmethod
-    def get_multiple(cls, parameters: List['Nucon']) -> Dict['Nucon', Union[float, int, bool, str, Enum]]:
-        return {param: cls.get(param) for param in parameters}
-
-    @classmethod
-    def get_all(cls) -> Dict['Nucon', Union[float, int, bool, str, Enum]]:
-        return cls.get_multiple(list(cls))
-
-    @classmethod
-    def set(cls, parameter: Union['Nucon', str], value: Union[float, int, bool, str, Enum], force: bool = False) -> None:
         if isinstance(parameter, str):
-            parameter = next((param for param in cls if param.id == parameter), None)
+            parameter = self._parameters[parameter]
-            if parameter is None:
+
-                raise ValueError(f"No parameter found with id '{parameter}'")
+        if self.dummy_mode:
+            return self._get_dummy_value(parameter)
+
+        value = self._query(parameter)
+        return self._parse_value(parameter, value)
+
+    def set(self, parameter: Union[str, NuconParameter], value: Union[float, int, bool, str, Enum], force: bool = False) -> None:
+        if isinstance(parameter, str):
+            parameter = self._parameters[parameter]
 
         if not force and not parameter.is_writable:
-            raise ValueError(f"Parameter {parameter.name} is not writable")
+            raise ValueError(f"Parameter {parameter} is not writable")
+        if not force and parameter.is_cheat and not self.cheat_mode:
+            raise ValueError(f"Parameter {parameter} is a cheat parameter. Enable cheat_mode on the Nucon instance or use force=True")
 
         if not force:
             parameter.check_in_range(value, raise_on_oob=True)
@@ -246,37 +473,198 @@ class Nucon(Enum):
         if parameter.enum_type and isinstance(value, parameter.enum_type):
             value = value.value
 
-        if NuconConfig.dummy_mode:
+        if self.dummy_mode:
-            print(f"Dummy mode: {'Force ' if force else ''}Setting {parameter.name} to {value}")
+            print(f"Dummy mode: {'Force ' if force else ''}Setting {parameter} to {value}")
             return
 
-        response = requests.post(NuconConfig.base_url, params={"variable": parameter.name, "value": str(value), "force": str(force).lower()})
+        self._set_value(parameter, str(value))
+
+    def get_type(self, parameter: Union[str, NuconParameter]) -> Type:
+        if isinstance(parameter, str):
+            parameter = self._parameters[parameter]
+        return parameter.param_type
+
+    def _query(self, parameter: NuconParameter) -> str:
+        response = requests.get(self.base_url, params={"variable": parameter.id})
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to query parameter {parameter.id}. Status code: {response.status_code}")
+
+        if response.text.strip() == 'NOT FOUND':
+            raise Exception(f"Failed to query parameter {parameter.id}. Returned 'NOT FOUND'")
+
+        return response.text.strip()
+
+    def _batch_query(self, param_names: List[str]) -> Dict[str, str]:
+        value = ','.join(param_names)
+        response = requests.get(self.base_url, params={"variable": "WEBSERVER_BATCH_GET", "value": value})
+
+        if response.status_code != 200:
+            raise Exception(f"Batch query failed. Status code: {response.status_code}")
+
+        data = response.json()
+        # Game returns {"values": {...}, "errors": {...}} wrapper
+        if isinstance(data, dict) and 'values' in data:
+            data = data['values']
+        return {k.upper(): ('' if v is None else str(v)) for k, v in data.items()}
+
+    def get_game_variable_names(self) -> List[str]:
+        response = requests.get(self.base_url, params={"variable": "WEBSERVER_LIST_VARIABLES"})
+
+        if response.status_code != 200:
+            raise Exception(f"Failed to get variable list. Status code: {response.status_code}")
+
+        names = []
+        for line in response.text.strip().splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            # Format is "GET:var1,var2,..." or "POST:var1,var2,..."
+            if ':' in line:
+                line = line.split(':', 1)[1]
+            names.extend(v.strip() for v in line.split(',') if v.strip())
+        # Return unique names, excluding known special (non-parameter) endpoints
+        seen = set()
+        result = []
+        for name in names:
+            upper = name.upper()
+            if upper not in seen and upper not in SPECIAL_VARIABLES:
+                seen.add(upper)
+                result.append(name)
+        return result
+
+    def _set_value(self, parameter: NuconParameter, value: str) -> None:
+        response = requests.post(self.base_url, params={"variable": parameter.id, "value": value})
         
         if response.status_code != 200:
-            raise Exception(f"Failed to set parameter {parameter.name}. Status code: {response.status_code}")
+            raise Exception(f"Failed to set parameter {parameter.id}. Status code: {response.status_code}")
 
-# Example usage
+    def _get_dummy_value(self, parameter: NuconParameter) -> Union[float, int, bool, str, Enum]:
-if __name__ == "__main__":
+        if parameter.enum_type:
-    # Enable dummy mode for testing
+            return next(iter(parameter.enum_type))
-    Nucon.set_dummy_mode(True)
+        elif parameter.param_type == float:
+            if parameter.max_val is not None and parameter.min_val is not None:
+                return (parameter.max_val - parameter.min_val) / 2 + parameter.min_val
+            else:
+                return 3.14
+        elif parameter.param_type == int:
+            if parameter.max_val is not None and parameter.min_val is not None:
+                return (parameter.max_val - parameter.min_val) // 2 + parameter.min_val
+            else:
+                return 42
+        elif parameter.param_type == bool:
+            return random.choice([True, False])
+        else:
+            return "dummy"
 
-    # Get a single parameter
+    def get_multiple_iter(self, parameters: List[Union[str, NuconParameter]]) -> Iterator[Tuple[str, Union[float, int, bool, str, Enum]]]:
-    core_temp = Nucon.CORE_TEMP.value
+        for param in parameters:
-    print(f"Core Temperature: {core_temp}")
+            if isinstance(param, str):
+                param_name = param
+                param_obj = self._parameters[param]
+            else:
+                param_name = next(name for name, p in self._parameters.items() if p is param)
+                param_obj = param
+            yield param_name, self.get(param_obj)
 
-    # Get a parameter with an enum
+    def get_multiple(self, parameters: List[Union[str, NuconParameter]]) -> Dict[str, Union[float, int, bool, str, Enum]]:
-    pump_status = Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value
+        return dict(self.get_multiple_iter(parameters))
-    print(f"Pump 0 Status: {pump_status}")
 
-    # Set a parameter with an enum
+    def get_all_iter(self) -> Iterator[Tuple[str, Union[float, int, bool, str, Enum]]]:
-    try:
+        if self.dummy_mode:
-        Nucon.GENERATOR_0_BREAKER.value = BreakerStatus.OPEN
+            yield from self.get_multiple_iter(self._parameters.keys())
-        print(f"Successfully set GENERATOR_0_BREAKER to {Nucon.GENERATOR_0_BREAKER.value}")
+            return
-    except ValueError as e:
-        print(f"Error: {e}")
 
-    # Get all parameters
+        try:
-    all_params = Nucon.get_all()
+            raw = self._batch_query(list(self._parameters.keys()))
-    print("All parameters:")
+        except Exception:
-    for param, value in all_params.items():
+            raw = {}
-        print(f"{param.name}: {value}")
+
+        for name, param in self._parameters.items():
+            if not param.is_readable:
+                continue
+            raw_value = raw.get(name)
+            # Batch query returns int codes for str-typed params; use individual query for those
+            if raw_value is not None and param.param_type != str:
+                yield name, self._parse_value(param, raw_value)
+            else:
+                yield name, self.get(param)
+
+    def get_all(self) -> Dict[str, Union[float, int, bool, str, Enum]]:
+        return dict(self.get_all_iter())
+
+    def get_all_readable(self) -> List[NuconParameter]:
+        return {name: param for name, param in self._parameters.items() if param.is_readable}
+
+    def get_all_writable(self) -> List[NuconParameter]:
+        return {name: param for name, param in self._parameters.items() if param.is_writable}
+
+    # --- Valve API ---
+    # Valves have a motorized actuator. OPEN/CLOSE power the motor toward that end-state;
+    # OFF cuts power and holds the current position. Normal resting state is OFF.
+    # The Value field (0-100) is the actual live position during travel.
+
+    def _post_valve_command(self, command: str, valve_name: str) -> None:
+        response = requests.post(self.base_url, params={"variable": command, "value": valve_name})
+        if response.status_code != 200:
+            raise Exception(f"Valve command {command} on '{valve_name}' failed. Status: {response.status_code}")
+
+    def get_valve(self, valve_name: str) -> Dict[str, Any]:
+        """Return current state dict for a single valve (from VALVE_PANEL_JSON)."""
+        valves = self.get_valves()
+        if valve_name not in valves:
+            raise KeyError(f"Valve '{valve_name}' not found")
+        return valves[valve_name]
+
+    def get_valves(self) -> Dict[str, Any]:
+        """Return state dict for all valves, keyed by valve name."""
+        response = requests.get(self.base_url, params={"variable": "VALVE_PANEL_JSON"})
+        if response.status_code != 200:
+            raise Exception(f"Failed to get valve panel. Status: {response.status_code}")
+        return response.json().get("valves", {})
+
+    def open_valve(self, valve_name: str) -> None:
+        """Power actuator toward open state. Send off_valve() once target is reached."""
+        self._post_valve_command("VALVE_OPEN", valve_name)
+
+    def close_valve(self, valve_name: str) -> None:
+        """Power actuator toward closed state. Send off_valve() once target is reached."""
+        self._post_valve_command("VALVE_CLOSE", valve_name)
+
+    def off_valve(self, valve_name: str) -> None:
+        """Cut actuator power, hold current position. Normal resting state."""
+        self._post_valve_command("VALVE_OFF", valve_name)
+
+    def open_valves(self, valve_names: List[str]) -> None:
+        for name in valve_names:
+            self.open_valve(name)
+
+    def close_valves(self, valve_names: List[str]) -> None:
+        for name in valve_names:
+            self.close_valve(name)
+
+    def off_valves(self, valve_names: List[str]) -> None:
+        for name in valve_names:
+            self.off_valve(name)
+
+    def set_dummy_mode(self, dummy_mode: bool) -> None:
+        self.dummy_mode = dummy_mode
+
+    def set_cheat_mode(self, cheat_mode: bool) -> None:
+        self.cheat_mode = cheat_mode
+
+    def __getattr__(self, name):
+        if isinstance(name, int):
+            return self.__getattr__(list(self._parameters.keys())[name])
+        if name in self._parameters:
+            return self._parameters[name]
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
+    def __getitem__(self, key):
+        return self.__getattr__(key)
+
+    def __dir__(self):
+        return list(super().__dir__()) + list(self._parameters.keys())
+
+    def __len__(self):
+        return sum(1 for p in self._parameters.values() if p.is_readable)

nucon/drake.py

@@ -0,0 +1,99 @@
+import os
+from pathlib import Path
+from PIL import Image, ImageDraw, ImageFont
+import textwrap
+
+# Get the directory where the script is located
+SCRIPT_DIR = Path(__file__).resolve().parent
+
+def get_asset_path(filename):
+    return str(SCRIPT_DIR / 'drake_assets' / filename)
+
+def fit_text_to_box(draw, text, font, max_width, max_height, line_spacing=1.2):
+    font_size = 1
+    while True:
+        font = ImageFont.truetype(font.path, font_size)
+        lines = textwrap.wrap(text, width=20)
+        line_height = font.getbbox("Ay")[3] - font.getbbox("Ay")[1]
+        total_height = line_height * len(lines) * line_spacing
+        max_line_width = max(font.getbbox(line)[2] - font.getbbox(line)[0] for line in lines)
+        
+        if max_line_width > max_width or total_height > max_height:
+            font_size -= 1
+            font = ImageFont.truetype(font.path, font_size)
+            break
+        font_size += 1
+    return font, lines
+
+def create_drake_meme(items):
+    # Load images
+    no_image = Image.open(get_asset_path('no.jpg'))
+    yes_image = Image.open(get_asset_path('yes.jpg'))
+    
+    # Set up meme dimensions
+    panel_width, panel_height = no_image.size
+    meme_width = panel_width * 2
+    meme_height = panel_height * len(items)
+    
+    # Create meme canvas
+    meme = Image.new("RGB", (meme_width, meme_height), "white")
+    
+    # Set up font (use a proper meme font)
+    font_path = get_asset_path('impact.ttf')
+    try:
+        base_font = ImageFont.truetype(font_path, 1)
+    except OSError:
+        # Fallback to default font if Impact is not available
+        base_font = ImageFont.load_default()
+    
+    for i, (is_yes, text) in enumerate(items):
+        # Paste the appropriate image
+        y_offset = i * panel_height
+        if is_yes:
+            meme.paste(yes_image, (0, y_offset))
+        else:
+            meme.paste(no_image, (0, y_offset))
+        
+        # Create text panel
+        text_panel = Image.new("RGB", (panel_width, panel_height), "white")
+        draw = ImageDraw.Draw(text_panel)
+        
+        # Fit and draw text
+        fitted_font, lines = fit_text_to_box(draw, text, base_font, panel_width - 20, panel_height - 20, line_spacing=1.2)
+        
+        line_height = fitted_font.getbbox("Ay")[3] - fitted_font.getbbox("Ay")[1]
+        total_height = line_height * len(lines) * 1.2
+        y_text = (panel_height - total_height) // 2
+        
+        for line in lines:
+            bbox = fitted_font.getbbox(line)
+            text_width = bbox[2] - bbox[0]
+            x_text = (panel_width - text_width) // 2
+            draw.text((x_text, y_text), line, font=fitted_font, fill="black")
+            y_text += line_height * 1.2
+        
+        # Paste text panel onto meme
+        meme.paste(text_panel, (panel_width, y_offset))
+
+    return meme
+
+default_items = [
+    (False, "Play Nucleares manually"),
+    (True, "Automate it with a script"),
+    (False, "But the web interface is tedious to use"),
+    (True, "Write an elegant libary to interface with the game and then use that to write the script"),
+    (False, "But I would still need to write the control policy by hand"),
+    (True, "Let's extend the libary such that it trains a policy via Reinforcement Learning"),
+    (False, "But RL takes a huge number of training samples"),
+    (True, "Extend the libary to also include an efficient simulator"),
+    (False, "But I don't know what the actual internal dynamics are"),
+    (True, "Extend the libary once more to also include a neural network dynamics model"),
+    (True, "And I'm gonna put a drake meme on the README"),
+    (False, "Online meme generators only support a single yes/no pair"),
+    (True, "Let's also add a drake meme generator to the libary"),
+    ]
+
+if __name__ == "__main__":
+    meme = create_drake_meme(default_items)
+    meme.save("drake_meme.jpg")
+    print("Meme saved as drake_meme.jpg")

nucon/model.py

@@ -0,0 +1,548 @@
+import numpy as np
+import time
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import random
+from enum import Enum
+from nucon import Nucon
+import pickle
+import os
+from typing import Union, Tuple, List, Dict
+
+Actors = {
+    'random': lambda nucon: lambda obs: {param.id: random.uniform(param.min_val, param.max_val) if param.min_val is not None and param.max_val is not None else 0 for param in nucon.get_all_writable().values()},
+    'null': lambda nucon: lambda obs: {},
+}
+
+# --- NN-based dynamics model ---
+
+class ReactorDynamicsNet(nn.Module):
+    def __init__(self, input_dim, output_dim, dropout=0.3):
+        super(ReactorDynamicsNet, self).__init__()
+        self.network = nn.Sequential(
+            nn.Linear(input_dim + 1, 128),  # +1 for time_delta
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, 128),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(128, output_dim)
+        )
+
+    def forward(self, state, time_delta):
+        x = torch.cat([state, time_delta], dim=-1)
+        return self.network(x)
+
+class ReactorDynamicsModel(nn.Module):
+    """
+    NN dynamics model predicting per-second rates of change (like ReactorKNNModel).
+
+    Inputs are z-score normalised; outputs are normalised rates.
+    forward() returns absolute next-state dict: cur + predicted_rate * time_delta.
+    forward_with_uncertainty() returns (next_state, 0.0) — no uncertainty estimate.
+    """
+    def __init__(self, input_params: List[str], output_params: List[str]):
+        super(ReactorDynamicsModel, self).__init__()
+        self.input_params = input_params
+        self.output_params = output_params
+        self.net = ReactorDynamicsNet(len(input_params), len(output_params))
+        # Normalisation stats set by fit()
+        self.register_buffer('_in_mean',  torch.zeros(len(input_params)))
+        self.register_buffer('_in_std',   torch.ones(len(input_params)))
+        self.register_buffer('_rate_mean', torch.zeros(len(output_params)))
+        self.register_buffer('_rate_std',  torch.ones(len(output_params)))
+
+    def fit_normalisation(self, dataset):
+        """Compute and store normalisation stats from a dataset."""
+        in_vecs, rate_vecs = [], []
+        for state, _action, next_state, dt in dataset:
+            if dt <= 0:
+                continue
+            in_vecs.append([state.get(p, 0.0) for p in self.input_params])
+            rate_vecs.append([(next_state.get(p, 0.0) - state.get(p, 0.0)) / dt
+                              for p in self.output_params])
+        ins   = np.array(in_vecs,   dtype=np.float32)
+        rates = np.array(rate_vecs, dtype=np.float32)
+        in_std = ins.std(0)
+        r_std  = rates.std(0)
+        self._in_mean.copy_(torch.from_numpy(ins.mean(0)))
+        self._in_std.copy_(torch.from_numpy(np.where(in_std < 1e-6, 1.0, in_std)))
+        self._rate_mean.copy_(torch.from_numpy(rates.mean(0)))
+        self._rate_std.copy_(torch.from_numpy(np.where(r_std < 1e-6, 1.0, r_std)))
+
+    def _normalise_input(self, t: torch.Tensor) -> torch.Tensor:
+        return (t - self._in_mean) / self._in_std
+
+    def _denormalise_rate(self, t: torch.Tensor) -> torch.Tensor:
+        return t * self._rate_std + self._rate_mean
+
+    def forward(self, state_dict, time_delta):
+        return self.forward_with_uncertainty(state_dict, time_delta)[0]
+
+    def forward_with_uncertainty(self, state_dict, time_delta, mc_samples=3):
+        """MC-Dropout uncertainty: run mc_samples stochastic forward passes.
+
+        Uncertainty is the mean normalised std across output dims, clipped to [0, 1].
+        0 = very confident (low variance), ~1 = high variance / OOD.
+        """
+        s = torch.tensor([state_dict.get(p, 0.0) for p in self.input_params],
+                         dtype=torch.float32).unsqueeze(0)
+        s_norm = self._normalise_input(s)
+        dt_t   = torch.tensor([[time_delta]], dtype=torch.float32)
+
+        # Keep dropout active for uncertainty sampling
+        self.net.train()
+        with torch.no_grad():
+            samples = torch.stack([self.net(s_norm, dt_t).squeeze(0)
+                                   for _ in range(mc_samples)])  # (mc_samples, out_dim)
+        self.net.eval()
+
+        rate_norm_mean = samples.mean(0)
+        rate_norm_std  = samples.std(0)
+
+        rate = self._denormalise_rate(rate_norm_mean)
+        cur  = torch.tensor([state_dict.get(p, 0.0) for p in self.output_params],
+                            dtype=torch.float32)
+        predicted = cur + rate * time_delta
+        pred_dict = {p: float(predicted[i]) for i, p in enumerate(self.output_params)}
+
+        # Uncertainty: mean coefficient of variation in normalised space, clipped to [0,1]
+        uncertainty = float(rate_norm_std.mean().clamp(0.0, 1.0))
+        return pred_dict, uncertainty
+
+# --- kNN-based dynamics model ---
+
+class ReactorKNNModel:
+    """
+    Non-parametric dynamics model using k-nearest neighbours.
+
+    For a query (state, game_delta):
+      1. Find the k dataset entries whose *state* is closest (L2 in normalised space).
+      2. For each neighbour compute the per-second rate-of-change:
+             rate_i = (next_state_i - state_i) / game_delta_i
+      3. Linearly scale to the requested game_delta:
+             predicted_delta_i = rate_i * game_delta
+      4. Return the inverse-distance-weighted average of those predicted deltas
+         added to the current output state.
+
+    The linear-in-time assumption means two datapoints at 0.5 s and 2 s contribute
+    equally once normalised by their own game_delta.
+    """
+
+    def __init__(self, input_params: List[str], output_params: List[str], k: int = 5):
+        self.input_params = input_params
+        self.output_params = output_params
+        self.k = k
+        self._states = None    # (n, d_in)  normalised state matrix
+        self._rates  = None    # (n, d_out) (next_out - cur_out) / game_delta
+        self._raw_states = None  # unnormalised, for mean/std computation
+        self._mean = None
+        self._std  = None
+
+    def fit(self, dataset):
+        """Build lookup tables from a collected dataset."""
+        raw, rates = [], []
+        for state, _action, next_state, game_delta in dataset:
+            if game_delta <= 0:
+                continue
+            s   = np.array([state[p]      for p in self.input_params],  dtype=np.float32)
+            cur = np.array([state[p]      for p in self.output_params], dtype=np.float32)
+            nxt = np.array([next_state[p] for p in self.output_params], dtype=np.float32)
+            raw.append(s)
+            rates.append((nxt - cur) / game_delta)
+
+        self._raw_states = np.array(raw)
+        self._rates      = np.array(rates)
+        self._mean = self._raw_states.mean(axis=0)
+        raw_std = self._raw_states.std(axis=0)
+        # Dimensions with zero variance in the training data carry no distance information.
+        # Use inf so they contribute 0 to normalised L2 (i.e., are ignored in kNN lookup).
+        self._std  = np.where(raw_std < 1e-6, np.inf, raw_std)
+        self._states = (self._raw_states - self._mean) / self._std
+
+    def _lookup(self, s: np.ndarray):
+        """Return (s_norm, idx, k) for the k nearest neighbours. s is a raw (d_in,) array."""
+        s_norm = (s - self._mean) / self._std
+        diff   = self._states - s_norm           # (n, d_in) broadcast
+        dists  = np.einsum('ij,ij->i', diff, diff)  # squared L2, faster than linalg.norm
+        k      = min(self.k, len(dists))
+        idx    = np.argpartition(dists, k - 1)[:k]
+        return s_norm, idx, k
+
+    def forward(self, state_dict: Dict, time_delta: float) -> Dict:
+        if self._states is None:
+            raise ValueError("Model not fitted. Call fit(dataset) first.")
+        return self.forward_with_uncertainty(state_dict, time_delta)[0]
+
+    def forward_with_uncertainty(self, state_dict: Dict, time_delta: float):
+        """Return (prediction_dict, uncertainty_scalar).
+
+        Uncertainty is the GP posterior std in normalised input space:
+          0 = query lies exactly on a training point (fully confident)
+          ~1 = query is far from all neighbours (maximally uncertain)
+        """
+        if self._states is None:
+            raise ValueError("Model not fitted. Call fit(dataset) first.")
+
+        s      = np.array([state_dict[p] for p in self.input_params], dtype=np.float32)
+        s_norm, idx, k = self._lookup(s)
+        X = self._states[idx]     # (k, d_in)
+        Y = self._rates[idx]      # (k, d_out)
+
+        # RBF kernel: k(a,b) = exp(-0.5 ||a-b||^2)
+        def rbf(A, B):
+            diff = A[:, None, :] - B[None, :, :]
+            return np.exp(-0.5 * np.einsum('ijk,ijk->ij', diff, diff))
+
+        K      = rbf(X, X) + 1e-4 * np.eye(k)
+        k_star = rbf(s_norm[None, :], X)[0]
+
+        K_inv      = np.linalg.inv(K)
+        mean_rates = k_star @ K_inv @ Y
+
+        var = max(0.0, 1.0 - float(k_star @ K_inv @ k_star))
+        std = float(np.sqrt(var))
+
+        cur_out   = np.array([state_dict[p] for p in self.output_params], dtype=np.float32)
+        predicted = cur_out + mean_rates * time_delta
+
+        pred_dict = {p: float(predicted[i]) for i, p in enumerate(self.output_params)}
+        return pred_dict, std
+
+# --- Mixture model ---
+
+class MixtureModel:
+    """Combines two dynamics models, selecting based on kNN uncertainty.
+
+    Uses knn_model when its uncertainty is below threshold (it's confident /
+    near training data). Falls back to nn_model when kNN is OOD.
+
+    Both models must implement forward_with_uncertainty(state_dict, time_delta).
+    input_params / output_params are taken from knn_model.
+    """
+    def __init__(self, knn_model, nn_model):
+        self.knn_model = knn_model
+        self.nn_model  = nn_model
+        self.input_params  = knn_model.input_params
+        self.output_params = knn_model.output_params
+
+    def forward(self, state_dict, time_delta):
+        return self.forward_with_uncertainty(state_dict, time_delta)[0]
+
+    def forward_with_uncertainty(self, state_dict, time_delta):
+        knn_pred, knn_u = self.knn_model.forward_with_uncertainty(state_dict, time_delta)
+        nn_pred,  nn_u  = self.nn_model.forward_with_uncertainty(state_dict, time_delta)
+        w_knn = 1.0 - knn_u   # high when kNN is confident
+        w_nn  = knn_u          # high when kNN is OOD
+        blended = {p: w_knn * knn_pred[p] + w_nn * nn_pred[p]
+                   for p in self.output_params}
+        uncertainty = w_knn * knn_u + w_nn * nn_u  # weighted uncertainty
+        return blended, uncertainty
+
+
+# --- Learner ---
+
+class NuconModelLearner:
+    def __init__(self, nucon=None, actor='null', dataset_path='nucon_dataset.pkl',
+                 time_delta: Union[float, Tuple[float, float]] = 1.0,
+                 include_valve_states: bool = False):
+        self.nucon = Nucon() if nucon is None else nucon
+        self.actor = Actors[actor](self.nucon) if actor in Actors else actor
+        self.dataset = self.load_dataset(dataset_path) or []
+        self.dataset_path = dataset_path
+        self.include_valve_states = include_valve_states
+        self.model = None
+        self.optimizer = None
+
+        # Exclude params with no physics signal
+        _JUNK_PARAMS = frozenset({'GAME_VERSION', 'TIME', 'TIME_STAMP', 'TIME_DAY',
+                                   'ALARMS_ACTIVE', 'FUN_IS_ENABLED', 'GAME_SIM_SPEED'})
+        candidate_params = {k: p for k, p in self.nucon.get_all_readable().items()
+                            if k not in _JUNK_PARAMS and p.param_type != str}
+        # Filter out params that return None (subsystem not installed).
+        # Retry until the game is reachable.
+        import requests as _requests
+        while True:
+            try:
+                test_state = {k: self.nucon.get(k) for k in candidate_params}
+                break
+            except (_requests.exceptions.ConnectionError,
+                    _requests.exceptions.Timeout):
+                print("Waiting for game to be reachable…")
+                time.sleep(5)
+        self.readable_params     = [k for k in candidate_params if test_state[k] is not None]
+        self.non_writable_params = [k for k in self.readable_params
+                                    if not self.nucon.get_all_readable()[k].is_writable]
+
+        # Optionally include valve positions (input only — valves are externally driven)
+        self.valve_keys = []
+        if include_valve_states:
+            valves = self.nucon.get_valves()
+            self.valve_keys = [f'VALVE__{name}' for name in sorted(valves.keys())]
+            self.readable_params = self.readable_params + self.valve_keys
+            # valve positions are input-only (not predicted as outputs)
+
+        if isinstance(time_delta, (int, float)):
+            self.time_delta = lambda: time_delta
+        elif isinstance(time_delta, tuple) and len(time_delta) == 2:
+            self.time_delta = lambda: random.uniform(*time_delta)
+        else:
+            raise ValueError("time_delta must be a float or a tuple of two floats")
+
+    def _get_state(self):
+        state = {}
+        for param_id in self.readable_params:
+            if param_id in self.valve_keys:
+                continue  # filled below
+            value = self.nucon.get(param_id)
+            if isinstance(value, Enum):
+                value = value.value
+            state[param_id] = value
+        if self.valve_keys:
+            valves = self.nucon.get_valves()
+            for key in self.valve_keys:
+                name = key[len('VALVE__'):]
+                state[key] = valves.get(name, {}).get('Value', 0.0)
+        return state
+
+    def collect_data(self, num_steps, save_every=10):
+        """
+        Collect state-transition tuples from the live game.
+
+        Sleeps wall_time = target_game_delta / sim_speed so that each stored
+        game_delta is uniform regardless of the game's simulation speed setting.
+
+        Saves the dataset every ``save_every`` steps so a crash doesn't lose
+        everything. On a connection error the step is skipped and collection
+        resumes once the game is reachable again (retries every 5 s).
+        """
+        import requests as _requests
+
+        def get_state_with_retry():
+            while True:
+                try:
+                    return self._get_state()
+                except (_requests.exceptions.ConnectionError,
+                        _requests.exceptions.Timeout) as e:
+                    print(f"Connection lost ({e}). Retrying in 5 s…")
+                    time.sleep(5)
+
+        state = get_state_with_retry()
+        collected = 0
+        for i in range(num_steps):
+            action = self.actor(state)
+            for param_id, value in action.items():
+                try:
+                    self.nucon.set(param_id, value)
+                except Exception:
+                    pass
+
+            target_game_delta = self.time_delta()
+            try:
+                sim_speed = self.nucon.GAME_SIM_SPEED.value or 1.0
+            except Exception:
+                sim_speed = 1.0
+            time.sleep(target_game_delta / sim_speed)
+
+            next_state = get_state_with_retry()
+            self.dataset.append((state, action, next_state, target_game_delta))
+            state = next_state
+            collected += 1
+
+            if collected % save_every == 0:
+                self.save_dataset()
+                print(f"  {collected}/{num_steps} steps collected, dataset saved.")
+
+        self.save_dataset()
+        print(f"Collection complete. {collected} steps, {len(self.dataset)} total samples.")
+
+    def train_model(self, batch_size=32, num_epochs=10, test_split=0.2, lr=1e-3):
+        """Train a neural-network dynamics model on the current dataset."""
+        if self.model is None:
+            self.model = ReactorDynamicsModel(self.readable_params, self.non_writable_params)
+        elif not isinstance(self.model, ReactorDynamicsModel):
+            raise ValueError("A kNN model is already loaded. Create a new learner to train an NN.")
+        self.model.fit_normalisation(self.dataset)
+        self.optimizer = optim.Adam(self.model.parameters(), lr=lr, weight_decay=1e-4)
+        random.shuffle(self.dataset)
+        split_idx = int(len(self.dataset) * (1 - test_split))
+        train_data = self.dataset[:split_idx]
+        test_data  = self.dataset[split_idx:]
+        for epoch in range(num_epochs):
+            train_loss = self._train_epoch(train_data, batch_size)
+            test_loss  = self._test_epoch(test_data)
+            print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}")
+
+    def fit_knn(self, k: int = 5):
+        """Fit a kNN/GP dynamics model from the current dataset (instantaneous, no gradient steps)."""
+        if self.model is None:
+            self.model = ReactorKNNModel(self.readable_params, self.non_writable_params, k=k)
+        elif not isinstance(self.model, ReactorKNNModel):
+            raise ValueError("An NN model is already loaded. Create a new learner to fit a kNN.")
+        self.model.fit(self.dataset)
+        print(f"kNN model fitted on {len(self.dataset)} samples.")
+
+    def predict_with_uncertainty(self, state_dict: Dict, time_delta: float):
+        """Return (prediction_dict, uncertainty_std). Only available after fit_knn()."""
+        if not isinstance(self.model, ReactorKNNModel):
+            raise ValueError("predict_with_uncertainty() requires a fitted kNN model (call fit_knn()).")
+        return self.model.forward_with_uncertainty(state_dict, time_delta)
+
+    def drop_well_fitted(self, error_threshold: float):
+        """Drop samples the current model already predicts well (MSE < threshold).
+
+        Keeps only hard/surprising transitions. Useful for NN training to focus
+        capacity on difficult regions of state space.
+        """
+        if self.model is None:
+            raise ValueError("No model fitted yet. Call train_model() or fit_knn() first.")
+        kept = []
+        for state, action, next_state, time_delta in self.dataset:
+            pred  = self.model.forward(state, time_delta)
+            error = sum((pred[p] - next_state[p]) ** 2 for p in self.non_writable_params)
+            if error > error_threshold:
+                kept.append((state, action, next_state, time_delta))
+        dropped = len(self.dataset) - len(kept)
+        self.dataset = kept
+        self.save_dataset()
+        print(f"drop_well_fitted: kept {len(kept)}, dropped {dropped} samples.")
+
+    def drop_redundant(self, min_state_distance: float, min_output_distance: float = 0.0):
+        """Drop near-duplicate samples, keeping only those that add coverage.
+
+        A sample is dropped only if *both* its input state and its output
+        transition are within the given distances of an already-kept sample
+        (L2 in z-scored space). If two samples share the same input state but
+        have different transitions they represent genuinely different dynamics
+        and are both kept regardless of `min_output_distance`.
+
+        Args:
+            min_state_distance:  minimum L2 distance in z-scored input space.
+            min_output_distance: minimum L2 distance in z-scored output-delta
+                                 space. Defaults to 0 (only input distance matters).
+        """
+        if not self.dataset:
+            return
+
+        in_params  = [p for p in self.readable_params if p not in self.valve_keys]
+        out_params = self.non_writable_params
+
+        all_states = np.array([[s[p]  for p in in_params]  for s, *_ in self.dataset], dtype=np.float32)
+        all_deltas = np.array([[ns[p] - s[p] for p in out_params]
+                               for s, _, ns, gd in self.dataset], dtype=np.float32)
+
+        s_mean, s_std = all_states.mean(0), all_states.std(0) + 1e-8
+        d_mean, d_std = all_deltas.mean(0), all_deltas.std(0) + 1e-8
+
+        s_norm = (all_states - s_mean) / s_std
+        d_norm = (all_deltas - d_mean) / d_std
+
+        kept_idx   = [0]
+        kept_s     = [s_norm[0]]
+        kept_d     = [d_norm[0]]
+
+        for i in range(1, len(self.dataset)):
+            s_dists = np.linalg.norm(np.array(kept_s) - s_norm[i], axis=1)
+            d_dists = np.linalg.norm(np.array(kept_d) - d_norm[i], axis=1)
+            # Drop only if close in BOTH spaces
+            if not np.any((s_dists < min_state_distance) & (d_dists < min_output_distance)):
+                kept_idx.append(i)
+                kept_s.append(s_norm[i])
+                kept_d.append(d_norm[i])
+
+        dropped = len(self.dataset) - len(kept_idx)
+        self.dataset = [self.dataset[i] for i in kept_idx]
+        self.save_dataset()
+        print(f"drop_redundant: kept {len(self.dataset)}, dropped {dropped} samples.")
+
+    def _train_epoch(self, data, batch_size):
+        self.model.train()
+        total_loss = 0
+        n_batches = 0
+        for i in range(0, len(data), batch_size):
+            batch = [s for s in data[i:i+batch_size] if s[3] > 0]
+            if not batch:
+                continue
+            states   = torch.tensor([[s[0].get(p, 0.0) for p in self.readable_params]  for s in batch], dtype=torch.float32)
+            targets  = torch.tensor([[(s[2].get(p, 0.0) - s[0].get(p, 0.0)) / s[3]    for p in self.non_writable_params] for s in batch], dtype=torch.float32)
+            dts      = torch.tensor([[s[3]] for s in batch], dtype=torch.float32)
+            s_norm   = self.model._normalise_input(states)
+            rate_norm_pred = self.model.net(s_norm, dts)
+            rate_norm_target = (targets - self.model._rate_mean) / self.model._rate_std
+            self.optimizer.zero_grad()
+            loss = torch.nn.functional.mse_loss(rate_norm_pred, rate_norm_target)
+            loss.backward()
+            self.optimizer.step()
+            total_loss += loss.item()
+            n_batches += 1
+        self.model.eval()
+        return total_loss / max(1, n_batches)
+
+    def _test_epoch(self, data):
+        total_loss = 0.0
+        n = 0
+        with torch.no_grad():
+            for state, _, next_state, dt in data:
+                if dt <= 0:
+                    continue
+                s_t    = torch.tensor([[state.get(p, 0.0) for p in self.readable_params]], dtype=torch.float32)
+                s_norm = self.model._normalise_input(s_t)
+                dt_t   = torch.tensor([[dt]], dtype=torch.float32)
+                rate_norm_pred   = self.model.net(s_norm, dt_t).squeeze(0)
+                target = torch.tensor([(next_state.get(p, 0.0) - state.get(p, 0.0)) / dt
+                                       for p in self.non_writable_params], dtype=torch.float32)
+                rate_norm_target = (target - self.model._rate_mean) / self.model._rate_std
+                total_loss += torch.nn.functional.mse_loss(rate_norm_pred, rate_norm_target).item()
+                n += 1
+        return total_loss / max(1, n)
+
+    def save_model(self, path):
+        if self.model is None:
+            raise ValueError("No model to save. Call train_model() or fit_knn() first.")
+        if isinstance(self.model, ReactorDynamicsModel):
+            torch.save({
+                'state_dict': self.model.state_dict(),
+                'input_params': self.model.input_params,
+                'output_params': self.model.output_params,
+            }, path)
+        else:
+            with open(path, 'wb') as f:
+                pickle.dump(self.model, f)
+
+    def load_model(self, path):
+        if path.endswith('.pkl'):
+            with open(path, 'rb') as f:
+                self.model = pickle.load(f)
+        else:
+            checkpoint = torch.load(path, weights_only=False)
+            if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+                m = ReactorDynamicsModel(checkpoint['input_params'], checkpoint['output_params'])
+                m.load_state_dict(checkpoint['state_dict'])
+                self.model = m
+            else:
+                # legacy plain state dict
+                self.model = ReactorDynamicsModel(self.readable_params, self.non_writable_params)
+                self.model.load_state_dict(checkpoint)
+
+    def save_dataset(self, path=None):
+        path = path or self.dataset_path
+        with open(path, 'wb') as f:
+            pickle.dump(self.dataset, f)
+
+    def load_dataset(self, path=None):
+        path = path or self.dataset_path
+        if os.path.exists(path):
+            with open(path, 'rb') as f:
+                return pickle.load(f)
+        return None
+
+    def merge_datasets(self, other_dataset_path):
+        other_dataset = self.load_dataset(other_dataset_path)
+        if not isinstance(other_dataset, list):
+            raise ValueError(
+                f"'{other_dataset_path}' does not contain a dataset (got {type(other_dataset).__name__}). "
+                f"Pass a dataset .pkl file, not a model file."
+            )
+        self.dataset.extend(other_dataset)
+        self.save_dataset()

nucon/rl.py

@@ -1,67 +1,169 @@
+import inspect
 import gymnasium as gym
 from gymnasium import spaces
 import numpy as np
 import time
-from typing import Dict, Any
+from typing import Dict, Any, Callable, List, Optional
-from .core import Nucon, BreakerStatus, PumpStatus, PumpDryStatus, PumpOverloadStatus
+from enum import Enum
+from nucon import Nucon, BreakerStatus, PumpStatus, PumpDryStatus, PumpOverloadStatus
+
+
+# ---------------------------------------------------------------------------
+# Reward / objective helpers
+# ---------------------------------------------------------------------------
+
+def _alarm_penalty(obs):
+    """Penalty proportional to number of active alarms. Only meaningful when running against the real game."""
+    raw = obs.get('ALARMS_ACTIVE', '')
+    if not raw or not raw.strip():
+        return 0.0
+    return -float(len(raw.split(',')))
 
 Objectives = {
-    "null": lambda obs: 0,
+    "null":          lambda obs: 0,
-    "coeff": lambda obj, coeff: lambda obs: obj(obs) * coeff,
+    "max_power":     lambda obs: obs["GENERATOR_0_KW"] + obs["GENERATOR_1_KW"] + obs["GENERATOR_2_KW"],
-
+    "episode_time":  lambda obs: obs["EPISODE_TIME"],
-    "max_power": lambda obs: obs["GENERATOR_0_KW"] + obs["GENERATOR_1_KW"] + obs["GENERATOR_2_KW"],
+    "alarm_penalty": _alarm_penalty,
-    "target_temperature": lambda goal_temp: lambda obs: (obs["CORE_TEMP"] - goal_temp) ** 2,
-    "episode_time": lambda obs: obs["EPISODE_TIME"],
 }
 
+def _uncertainty_penalty(start=0.3, scale=1.0, mode='l2'):
+    excess = lambda obs: max(0.0, obs.get('SIM_UNCERTAINTY', 0.0) - start)
+    if mode == 'l2':
+        return lambda obs: -scale * excess(obs) ** 2
+    elif mode == 'linear':
+        return lambda obs: -scale * excess(obs)
+    else:
+        raise ValueError(f"Unknown mode '{mode}'. Use 'l2' or 'linear'.")
+
+def _uncertainty_abort(threshold=0.7):
+    return lambda obs: 1.0 if obs.get('SIM_UNCERTAINTY', 0.0) >= threshold else 0.0
+
+Parameterized_Objectives = {
+    "target_temperature":  lambda goal_temp: lambda obs: -((obs["CORE_TEMP"] - goal_temp) ** 2),
+    "target_gap":          lambda goal_gap:  lambda obs: -((obs["CORE_TEMP"] - obs["CORE_TEMP_MIN"] - goal_gap) ** 2),
+    "temp_below":          lambda max_temp:  lambda obs: -(np.clip(obs["CORE_TEMP"] - max_temp,  0, np.inf) ** 2),
+    "temp_above":          lambda min_temp:  lambda obs: -(np.clip(min_temp - obs["CORE_TEMP"],  0, np.inf) ** 2),
+    "temp_below_linear":   lambda max_temp:  lambda obs: -np.clip(obs["CORE_TEMP"] - max_temp,  0, np.inf),
+    "temp_above_linear":   lambda min_temp:  lambda obs: -np.clip(min_temp - obs["CORE_TEMP"],  0, np.inf),
+    "constant":            lambda constant:  lambda obs: constant,
+    "uncertainty_penalty": _uncertainty_penalty,  # (start, scale, mode) -> (obs) -> float
+}
+
+Parameterized_Terminators = {
+    "uncertainty_abort": _uncertainty_abort,  # (threshold,) -> (obs) -> float
+}
+
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+def _build_flat_action_space(nucon, obs_param_set=None, delta_action_scale=None):
+    """Return (Box, ordered_param_ids, param_ranges).
+
+    If delta_action_scale is set, the action space is [-1, 1]^n and actions are
+    treated as normalised deltas: actual_delta = action * delta_action_scale * (max - min).
+    Otherwise the action space spans [min_val, max_val] per param (absolute values).
+    """
+    params = []
+    lows, highs, ranges = [], [], []
+    for param_id, param in nucon.get_all_writable().items():
+        if not param.is_readable or param.is_cheat:
+            continue
+        if obs_param_set is not None and param_id not in obs_param_set:
+            continue
+        if param.min_val is None or param.max_val is None:
+            continue  # SAC requires finite action bounds
+        sp = _build_param_space(param)
+        if sp is None:
+            continue
+        params.append(param_id)
+        lows.append(sp.low[0])
+        highs.append(sp.high[0])
+        ranges.append(sp.high[0] - sp.low[0])
+    if delta_action_scale is not None:
+        n = len(params)
+        box = spaces.Box(low=-np.ones(n, dtype=np.float32),
+                         high=np.ones(n, dtype=np.float32), dtype=np.float32)
+    else:
+        box = spaces.Box(low=np.array(lows, dtype=np.float32),
+                         high=np.array(highs, dtype=np.float32), dtype=np.float32)
+    return box, params, np.array(lows, dtype=np.float32), np.array(ranges, dtype=np.float32)
+
+
+def _unflatten_action(flat_action, param_ids):
+    return {pid: float(flat_action[i]) for i, pid in enumerate(param_ids)}
+
+
+def _build_param_space(param):
+    """Return a gymnasium Box for a single NuconParameter, or None if unsupported."""
+    if param.param_type in (float, int):
+        lo = param.min_val if param.min_val is not None else -np.inf
+        hi = param.max_val if param.max_val is not None else np.inf
+        return spaces.Box(low=lo, high=hi, shape=(1,), dtype=np.float32)
+    elif param.param_type == bool:
+        return spaces.Box(low=0, high=1, shape=(1,), dtype=np.float32)
+    elif param.param_type == str:
+        return None
+    elif issubclass(param.param_type, Enum):
+        return spaces.Box(low=0, high=len(param.param_type) - 1, shape=(1,), dtype=np.float32)
+    return None
+
+
+def _apply_action(nucon, action):
+    for param_id, value in action.items():
+        param = nucon._parameters[param_id]
+        v = float(np.asarray(value).flat[0])
+        if param.param_type == bool:
+            value = v >= 0.5  # [0,1] space: above midpoint → True
+        elif issubclass(param.param_type, Enum):
+            value = param.param_type(int(v))
+        else:
+            value = param.param_type(v)
+            if param.min_val is not None and param.max_val is not None:
+                value = param.param_type(np.clip(value, param.min_val, param.max_val))
+        nucon.set(param, value)
+
+
+# ---------------------------------------------------------------------------
+# NuconEnv
+# ---------------------------------------------------------------------------
+
 class NuconEnv(gym.Env):
     metadata = {'render_modes': ['human']}
 
-    def __init__(self, render_mode=None, seconds_per_step=5, objectives=['null'], terminators=['null'], terminate_above=0):
+    def __init__(self, nucon=None, simulator=None, render_mode=None, seconds_per_step=5,
+                 objectives=['null'], terminators=['null'], objective_weights=None, terminate_above=0):
         super().__init__()
 
         self.render_mode = render_mode
         self.seconds_per_step = seconds_per_step
-        self.terminate_at = terminate_at
+        if objective_weights is None:
-        
+            objective_weights = [1.0 for _ in objectives]
-        # Define observation space
+        self.objective_weights = objective_weights
-        obs_spaces = {'EPISODE_TIME': spaces.Box(low=0, high=np.inf, shape=(1,), dtype=np.float32)}
+        self.terminate_above = terminate_above
-        for param in Nucon.get_all_readable():
+        self.simulator = simulator
-            if param.param_type == float:
-                obs_spaces[param.id] = spaces.Box(low=param.min_val or -np.inf, high=param.max_val or np.inf, shape=(1,), dtype=np.float32)
-            elif param.param_type == int:
-                if param.min_val is not None and param.max_val is not None:
-                    obs_spaces[param.id] = spaces.Box(low=param.min_val, high=param.max_val, shape=(1,), dtype=np.float32)
-                else:
-                    obs_spaces[param.id] = spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32)
-            elif param.param_type == bool:
-                obs_spaces[param.id] = spaces.Box(low=0, high=1, shape=(1,), dtype=np.float32)
-            elif issubclass(param.param_type, Enum):
-                obs_spaces[param.id] = spaces.Box(low=0, high=1, shape=(len(param.param_type),), dtype=np.float32)
-            else:
-                raise ValueError(f"Unsupported observation parameter type: {param.param_type}")
 
+        if nucon is None:
+            nucon = Nucon(port=simulator.port) if simulator else Nucon()
+        self.nucon = nucon
+
+        # Observation space — SIM_UNCERTAINTY included when a simulator is present
+        obs_spaces = {'EPISODE_TIME': spaces.Box(low=0, high=np.inf, shape=(1,), dtype=np.float32)}
+        if simulator is not None:
+            obs_spaces['SIM_UNCERTAINTY'] = spaces.Box(low=0.0, high=1.0, shape=(1,), dtype=np.float32)
+        for param_id, param in self.nucon.get_all_readable().items():
+            sp = _build_param_space(param)
+            if sp is not None:
+                obs_spaces[param_id] = sp
         self.observation_space = spaces.Dict(obs_spaces)
 
-        # Define action space
+        self.action_space, self._action_params, self._action_lows, self._action_ranges = \
-        action_spaces = {}
+            _build_flat_action_space(self.nucon)
-        for param in Nucon.get_all_writable():
-            if param.param_type == float:
-                action_spaces[param.id] = spaces.Box(low=param.min_val or -np.inf, high=param.max_val or np.inf, shape=(1,), dtype=np.float32)
-            elif param.param_type == int:
-                if param.min_val is not None and param.max_val is not None:
-                    action_spaces[param.id] = spaces.Box(low=param.min_val, high=param.max_val, shape=(1,), dtype=np.float32)
-                else:
-                    action_spaces[param.id] = spaces.Box(low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32)
-            elif param.param_type == bool:
-                action_spaces[param.id] = spaces.Box(low=0, high=1, shape=(1,), dtype=np.float32)
-            elif issubclass(param.param_type, Enum):
-                action_spaces[param.id] = spaces.Box(low=0, high=1, shape=(len(param.param_type),), dtype=np.float32)
-            else:
-                raise ValueError(f"Unsupported action parameter type: {param.param_type}")
-
-        self.action_space = spaces.Dict(action_spaces)
 
+        self.objectives = []
+        self.terminators = []
         for objective in objectives:
             if objective in Objectives:
                 self.objectives.append(Objectives[objective])
@@ -69,7 +171,6 @@ class NuconEnv(gym.Env):
                 self.objectives.append(objective)
             else:
                 raise ValueError(f"Unsupported objective: {objective}")
-
         for terminator in terminators:
             if terminator in Objectives:
                 self.terminators.append(Objectives[terminator])
@@ -78,66 +179,388 @@ class NuconEnv(gym.Env):
             else:
                 raise ValueError(f"Unsupported terminator: {terminator}")
 
-    def _get_obs(self):
+    def _get_obs(self, sim_uncertainty=None):
         obs = {}
-        for param in Nucon.get_all_readable():
+        for param_id, param in self.nucon.get_all_readable().items():
-            value = Nucon.get(param)
+            if param.param_type == str or param_id not in self.observation_space.spaces:
+                continue
+            value = self.nucon.get(param_id)
             if isinstance(value, Enum):
                 value = value.value
-            obs[param.id] = value
+            obs[param_id] = value
-        obs["EPISODE_TIME"] = self._total_steps * self.seconds_per_step
+        obs['EPISODE_TIME'] = self._total_steps * self.seconds_per_step
+        if 'SIM_UNCERTAINTY' in self.observation_space.spaces:
+            obs['SIM_UNCERTAINTY'] = sim_uncertainty if sim_uncertainty is not None else 0.0
         return obs
 
-    def _get_info(self):
+    def _get_info(self, obs):
-        info = {'objectives': {}}
+        info = {'objectives': {}, 'objectives_weighted': {}}
-        for objective in self.objectives:
+        for objective, weight in zip(self.objectives, self.objective_weights):
-            info['objectives'][objective.__name__] = objective(self._get_obs())
+            obj = objective(obs)
+            name = getattr(objective, '__name__', repr(objective))
+            info['objectives'][name] = obj
+            info['objectives_weighted'][name] = obj * weight
         return info
-    
+
     def reset(self, seed=None, options=None):
         super().reset(seed=seed)
-
         self._total_steps = 0
         observation = self._get_obs()
-        info = self._get_info()
+        return observation, self._get_info(observation)
-
-        return observation, info
 
     def step(self, action):
-        # Apply the action to the Nucon system
+        _apply_action(self.nucon, _unflatten_action(action, self._action_params))
-        for param_id, value in action.items():
-            param = next(p for p in Nucon if p.id == param_id)
-            if issubclass(param.param_type, Enum):
-                value = param.param_type(value)
-            if param.min_val is not None and param.max_val is not None:
-                value = np.clip(value, param.min_val, param.max_val)
-            Nucon.set(param, value)
 
-        observation = self._get_obs()
+        # Advance sim (or sleep) — get uncertainty for obs injection
-        terminated = np.sum([terminator(observation) for terminator in self.terminators]) > self.terminate_above
         truncated = False
-        info = self._get_info()
+        uncertainty = None
-        reward = sum(obj for obj in info['objectives'].values())
+        if self.simulator:
+            uncertainty = self.simulator.update(self.seconds_per_step, return_uncertainty=True)
+        else:
+            sim_speed = self.nucon.GAME_SIM_SPEED.value or 1.0
+            time.sleep(self.seconds_per_step / sim_speed)
 
         self._total_steps += 1
-        time.sleep(self.seconds_per_step)
+        observation = self._get_obs(sim_uncertainty=uncertainty)
+        info = self._get_info(observation)
+        reward = sum(obj for obj in info['objectives_weighted'].values())
+        terminated = np.sum([t(observation) for t in self.terminators]) > self.terminate_above
         return observation, reward, terminated, truncated, info
 
     def render(self):
-        if self.render_mode == "human":
+        pass
-            pass
 
     def close(self):
         pass
 
-    def _flatten_action(self, action):
-        return np.concatenate([v.flatten() for v in action.values()])
-
-    def _unflatten_action(self, flat_action):
-        return {k: v.reshape(1, -1) for k, v in self.action_space.items()}
-
     def _flatten_observation(self, observation):
-        return np.concatenate([v.flatten() for v in observation.values()])
+        return np.concatenate([np.asarray(v).flatten() for v in observation.values()])
 
-    def _unflatten_observation(self, flat_observation):
+
-        return {k: v.reshape(1, -1) for k, v in self.observation_space.items()}
+# ---------------------------------------------------------------------------
+# NuconGoalEnv
+# ---------------------------------------------------------------------------
+
+class NuconGoalEnv(gym.Env):
+    """
+    Goal-conditioned reactor environment compatible with SB3 HER (Hindsight Experience Replay).
+
+    Observation is a Dict with three keys:
+      - 'observation':   all readable non-goal, non-str params + SIM_UNCERTAINTY (when sim active)
+      - 'achieved_goal': current values of goal_params, normalised to [0, 1] within goal_range
+      - 'desired_goal':  target values sampled each episode, normalised to [0, 1]
+
+    ``SIM_UNCERTAINTY`` in 'observation' lets reward_fn / terminators reference uncertainty directly.
+
+    reward_fn signature: ``(achieved, desired)`` or ``(achieved, desired, obs)`` — the 3-arg form
+    receives the full observation dict (including SIM_UNCERTAINTY) for uncertainty-aware shaping.
+
+    Usage with SB3 HER::
+
+        from stable_baselines3 import SAC
+        from stable_baselines3.common.buffers import HerReplayBuffer
+        from nucon.rl import NuconGoalEnv, UncertaintyPenalty, UncertaintyAbort
+
+        env = NuconGoalEnv(
+            goal_params=['GENERATOR_0_KW', 'GENERATOR_1_KW', 'GENERATOR_2_KW'],
+            goal_range={'GENERATOR_0_KW': (0, 1200), 'GENERATOR_1_KW': (0, 1200), 'GENERATOR_2_KW': (0, 1200)},
+            tolerance=0.05,
+            simulator=simulator,
+            # uncertainty-aware reward: penalise OOD, abort if too far out
+            reward_fn=lambda ag, dg, obs: (
+                -(np.linalg.norm(ag - dg) ** 2)
+                - 2.0 * max(0, obs.get('SIM_UNCERTAINTY', 0) - 0.3) ** 2
+            ),
+            terminators=[UncertaintyAbort(threshold=0.7)],
+        )
+        model = SAC('MultiInputPolicy', env, replay_buffer_class=HerReplayBuffer)
+        model.learn(total_timesteps=500_000)
+    """
+
+    metadata = {'render_modes': ['human']}
+
+    def __init__(
+        self,
+        goal_params,
+        goal_range=None,
+        reward_fn=None,
+        tolerance=None,
+        nucon=None,
+        simulator=None,
+        render_mode=None,
+        seconds_per_step=5,
+        terminators=None,
+        terminate_above=0,
+        additional_objectives=None,
+        additional_objective_weights=None,
+        obs_params=None,
+        action_params=None,
+        init_states=None,
+        delta_action_scale=None,
+        goal_sampling_std=None,
+    ):
+        super().__init__()
+
+        self.render_mode = render_mode
+        self.seconds_per_step = seconds_per_step
+        self._delta_action_scale = delta_action_scale
+        self.terminate_above = terminate_above
+        self.simulator = simulator
+        self.goal_params = list(goal_params)
+        self.tolerance = tolerance
+
+        if nucon is None:
+            nucon = Nucon(port=simulator.port) if simulator else Nucon()
+        self.nucon = nucon
+
+        all_readable = self.nucon.get_all_readable()
+        for pid in self.goal_params:
+            if pid not in all_readable:
+                raise ValueError(f"Goal param '{pid}' is not a readable parameter")
+
+        goal_range = goal_range or {}
+        self._goal_low = np.array([
+            goal_range.get(pid, (all_readable[pid].min_val or 0.0, all_readable[pid].max_val or 1.0))[0]
+            for pid in self.goal_params
+        ], dtype=np.float32)
+        self._goal_high = np.array([
+            goal_range.get(pid, (all_readable[pid].min_val or 0.0, all_readable[pid].max_val or 1.0))[1]
+            for pid in self.goal_params
+        ], dtype=np.float32)
+        self._goal_range = self._goal_high - self._goal_low
+        self._goal_range[self._goal_range == 0] = 1.0
+
+        # Detect reward_fn arity for backward compat (2-arg vs 3-arg)
+        self._reward_fn = reward_fn
+        if reward_fn is not None:
+            n_args = len(inspect.signature(reward_fn).parameters)
+            self._reward_fn_wants_obs = n_args >= 3
+        else:
+            self._reward_fn_wants_obs = False
+
+        # Observation params: model.input_params defines the canonical list — the same set is
+        # used whether training in sim or deploying to the real game (the game simply has more
+        # params available; we query only the subset we care about).
+        # Explicit obs_params overrides everything (use when deploying to real game without sim).
+        # SB3 HER requires observation to be a flat Box, not a nested Dict.
+        goal_set = set(self.goal_params)
+        self._obs_with_uncertainty = simulator is not None
+        if obs_params is not None:
+            base_params = [p for p in obs_params if p not in goal_set]
+        elif simulator is not None and hasattr(simulator, 'model') and simulator.model is not None:
+            base_params = [p for p in simulator.model.input_params
+                           if p not in goal_set and p in all_readable
+                           and _build_param_space(all_readable[p]) is not None]
+        else:
+            base_params = [p for p, param in all_readable.items()
+                           if p not in goal_set and _build_param_space(param) is not None]
+        # SIM_UNCERTAINTY is not in _obs_params — it's not available at deployment on the real game
+        self._obs_params = base_params
+
+        n_goals = len(self.goal_params)
+        self.observation_space = spaces.Dict({
+            'observation':   spaces.Box(low=-np.inf, high=np.inf,
+                                        shape=(len(self._obs_params),), dtype=np.float32),
+            'achieved_goal': spaces.Box(low=0.0, high=1.0, shape=(n_goals,), dtype=np.float32),
+            'desired_goal':  spaces.Box(low=0.0, high=1.0, shape=(n_goals,), dtype=np.float32),
+        })
+
+        # Action space: writable params within the obs param set, or an explicit override list.
+        action_set = set(action_params) if action_params is not None else set(base_params)
+        self.action_space, self._action_params, self._action_lows, self._action_ranges = \
+            _build_flat_action_space(self.nucon, action_set, delta_action_scale)
+
+        self._terminators = terminators or []
+        _objs = additional_objectives or []
+        self._objectives = [Objectives[o] if isinstance(o, str) else o for o in _objs]
+        self._objective_weights = additional_objective_weights or [1.0] * len(self._objectives)
+        self._init_states = init_states  # list of state dicts to sample on reset
+        self._goal_sampling_std = goal_sampling_std  # Gaussian std in normalised goal space; None → uniform
+        self._desired_goal = np.zeros(n_goals, dtype=np.float32)
+        self._total_steps = 0
+
+    def compute_reward(self, achieved_goal, desired_goal, info):
+        """Dense negative L2, sparse with tolerance, or custom reward_fn."""
+        obs_named = info.get('obs_named', {}) if isinstance(info, dict) else {}
+        if self._reward_fn is not None:
+            if self._reward_fn_wants_obs:
+                return self._reward_fn(achieved_goal, desired_goal, obs_named)
+            return self._reward_fn(achieved_goal, desired_goal)
+        dist = np.linalg.norm(achieved_goal - desired_goal, axis=-1)
+        if self.tolerance is not None:
+            return (dist <= self.tolerance).astype(np.float32) - 1.0
+        return -dist
+
+    def _read_goal_values(self):
+        raw = np.array([self.nucon.get(pid) or 0.0 for pid in self.goal_params], dtype=np.float32)
+        return np.clip((raw - self._goal_low) / self._goal_range, 0.0, 1.0)
+
+    def _read_obs(self, sim_uncertainty=None):
+        """Return (gym_obs_dict, reward_obs_dict).
+
+        When a simulator is attached, reads directly from sim.parameters (no HTTP).
+        Otherwise falls back to a single batch HTTP request.
+        """
+        def _to_float(v):
+            if v is None:
+                return 0.0
+            return float(v.value if isinstance(v, Enum) else v)
+
+        if self.simulator is not None:
+            # Direct in-process read — no HTTP overhead
+            def _get(pid):
+                return _to_float(self.simulator.get(pid))
+        else:
+            raw = self.nucon._batch_query(self._obs_params + self.goal_params)
+            all_params = self.nucon.get_all_readable()
+            def _get(pid):
+                try:
+                    v = self.nucon._parse_value(all_params[pid], raw.get(pid, '0'))
+                    return _to_float(v)
+                except Exception:
+                    return 0.0
+
+        reward_obs = {}
+        if self._obs_with_uncertainty:
+            reward_obs['SIM_UNCERTAINTY'] = float(sim_uncertainty) if sim_uncertainty is not None else 0.0
+        for pid in self._obs_params:
+            reward_obs[pid] = _get(pid)
+
+        obs_vec  = np.array([reward_obs[p] for p in self._obs_params], dtype=np.float32)
+        goal_raw = np.array([_get(p) for p in self.goal_params], dtype=np.float32)
+        achieved = np.clip((goal_raw - self._goal_low) / self._goal_range, 0.0, 1.0)
+        gym_obs  = {'observation': obs_vec, 'achieved_goal': achieved,
+                    'desired_goal': self._desired_goal.copy()}
+        return gym_obs, reward_obs
+
+    def reset(self, seed=None, options=None):
+        super().reset(seed=seed)
+        self._total_steps = 0
+        rng = np.random.default_rng(seed)
+        if self._init_states is not None and self.simulator is not None:
+            state = self._init_states[rng.integers(len(self._init_states))]
+            for k, v in state.items():
+                try:
+                    self.simulator.set(k, v, force=True)
+                except Exception:
+                    pass
+        if self._goal_sampling_std is not None:
+            # Sample goal as Gaussian delta from current state — usually a small change,
+            # occasionally a large one.
+            current = np.array([
+                float(self.simulator.get(p) if self.simulator else 0.0)
+                for p in self.goal_params
+            ], dtype=np.float32)
+            current_norm = np.clip((current - self._goal_low) / self._goal_range, 0.0, 1.0)
+            delta = rng.normal(0.0, self._goal_sampling_std, size=len(self.goal_params))
+            self._desired_goal = np.clip(current_norm + delta, 0.0, 1.0).astype(np.float32)
+        else:
+            self._desired_goal = rng.uniform(0.0, 1.0, size=len(self.goal_params)).astype(np.float32)
+        gym_obs, _ = self._read_obs()
+        return gym_obs, {}
+
+    def step(self, action):
+        flat = np.asarray(action, dtype=np.float32)
+        if self._delta_action_scale is not None:
+            # Compute absolute values from deltas, reading current state
+            if self.simulator is None:
+                raw_current = self.nucon._batch_query(self._action_params)
+                all_params = self.nucon.get_all_readable()
+            absolute = {}
+            for i, pid in enumerate(self._action_params):
+                param = self.nucon._parameters[pid]
+                if param.param_type == bool:
+                    absolute[pid] = 1.0 if flat[i] > 0 else 0.0
+                else:
+                    if self.simulator is not None:
+                        v = self.simulator.get(pid)
+                        current = float(v.value if isinstance(v, Enum) else v) if v is not None else 0.0
+                    else:
+                        try:
+                            v = self.nucon._parse_value(all_params[pid], raw_current.get(pid, '0'))
+                            current = float(v.value if isinstance(v, Enum) else v)
+                        except Exception:
+                            current = 0.0
+                    delta = float(flat[i]) * self._delta_action_scale * self._action_ranges[i]
+                    absolute[pid] = float(np.clip(current + delta,
+                                                  self._action_lows[i],
+                                                  self._action_lows[i] + self._action_ranges[i]))
+        else:
+            absolute = _unflatten_action(flat, self._action_params)
+
+        if self.simulator is not None:
+            # Write directly to sim — skip HTTP entirely
+            for pid, val in absolute.items():
+                try:
+                    self.simulator.set(pid, val, force=True)
+                except Exception:
+                    pass
+        else:
+            _apply_action(self.nucon, absolute)
+
+        if self.simulator:
+            uncertainty = self.simulator.update(self.seconds_per_step, return_uncertainty=True)
+        else:
+            sim_speed = self.nucon.GAME_SIM_SPEED.value or 1.0
+            time.sleep(self.seconds_per_step / sim_speed)
+            uncertainty = None
+
+        self._total_steps += 1
+        gym_obs, reward_obs = self._read_obs(sim_uncertainty=uncertainty)
+        info = {'achieved_goal': gym_obs['achieved_goal'], 'desired_goal': gym_obs['desired_goal'],
+                'obs_named': reward_obs}
+        reward  = float(self.compute_reward(gym_obs['achieved_goal'], gym_obs['desired_goal'], info))
+        reward += sum(w * o(reward_obs) for o, w in zip(self._objectives, self._objective_weights))
+        terminated = any(t(reward_obs) > self.terminate_above for t in self._terminators)
+        return gym_obs, reward, terminated, False, info
+
+    def render(self):
+        pass
+
+    def close(self):
+        pass
+
+
+# ---------------------------------------------------------------------------
+# Registration
+# ---------------------------------------------------------------------------
+
+def register_nucon_envs():
+    gym.register(
+        id='Nucon-max_power-v0',
+        entry_point='nucon.rl:NuconEnv',
+        kwargs={'seconds_per_step': 5, 'objectives': ['max_power']}
+    )
+    gym.register(
+        id='Nucon-target_temperature_350-v0',
+        entry_point='nucon.rl:NuconEnv',
+        kwargs={'seconds_per_step': 5, 'objectives': [Parameterized_Objectives['target_temperature'](goal_temp=350)]}
+    )
+    gym.register(
+        id='Nucon-safe_max_power-v0',
+        entry_point='nucon.rl:NuconEnv',
+        kwargs={'seconds_per_step': 5,
+                'objectives': [Parameterized_Objectives['temp_above'](min_temp=310),
+                               Parameterized_Objectives['temp_below'](max_temp=365), 'max_power'],
+                'objective_weights': [1, 10, 1/100_000]}
+    )
+    gym.register(
+        id='Nucon-goal_power-v0',
+        entry_point='nucon.rl:NuconGoalEnv',
+        kwargs={
+            'goal_params': ['GENERATOR_0_KW', 'GENERATOR_1_KW', 'GENERATOR_2_KW'],
+            'goal_range': {'GENERATOR_0_KW': (0.0, 1200.0), 'GENERATOR_1_KW': (0.0, 1200.0), 'GENERATOR_2_KW': (0.0, 1200.0)},
+            'seconds_per_step': 5,
+        }
+    )
+    gym.register(
+        id='Nucon-goal_temp-v0',
+        entry_point='nucon.rl:NuconGoalEnv',
+        kwargs={
+            'goal_params': ['CORE_TEMP'],
+            'goal_range': {'CORE_TEMP': (280.0, 380.0)},
+            'seconds_per_step': 5,
+        }
+    )
+
+register_nucon_envs()

nucon/sim.py

@@ -0,0 +1,356 @@
+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 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()

pyproject.toml

@@ -25,7 +25,17 @@ dependencies = [
 ]
 
 [project.urls]
-Homepage = ""
+Homepage = "https://git.dominik-roth.eu/dodox/nucon"
 
 [project.optional-dependencies]
-dev = ["pytest"]
+dev = ["pytest", "flask", "numpy"]
+rl = ["gymnasium", "numpy"]
+sim = ["torch", "flask"]
+model = ["torch", "numpy"]
+drake = ["Pillow"]
+
+[tool.setuptools.package-data]
+"nucon" = ["drake_assets/*"]
+
+[tool.setuptools]
+include-package-data = true

scripts/collect_dataset.py

@@ -0,0 +1,55 @@
+"""Collect a dynamics dataset from the running Nucleares game.
+
+Play the game normally while this script runs in the background.
+It records state transitions every `time_delta` game-seconds and
+saves them incrementally so nothing is lost if you quit early.
+
+Usage:
+    python scripts/collect_dataset.py                          # default settings
+    python scripts/collect_dataset.py --steps 2000 --delta 5  # faster sampling
+    python scripts/collect_dataset.py --out my_dataset.pkl
+
+The saved dataset is a list of (state_before, action_dict, state_after, time_delta)
+tuples compatible with NuconModelLearner.fit_knn() and train_model().
+
+Tips for good data:
+  - Cover a range of operating states: startup, ramp, steady-state, shutdown.
+  - Vary individual rod bank positions, pump speeds, and MSCV setpoints.
+  - Collect at least 500 samples for kNN-GP; 5000+ for the NN backend.
+  - Merge multiple sessions with NuconModelLearner.merge_datasets().
+"""
+import argparse
+import pickle
+from nucon.model import NuconModelLearner
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--steps', type=int, default=1000,
+                    help='Number of samples to collect (default: 1000)')
+parser.add_argument('--delta', type=float, default=10.0,
+                    help='Game-seconds between samples (default: 10.0)')
+parser.add_argument('--out', default='reactor_dataset.pkl',
+                    help='Output path for dataset (default: reactor_dataset.pkl)')
+parser.add_argument('--merge', default=None,
+                    help='Existing dataset to merge into before saving')
+args = parser.parse_args()
+
+learner = NuconModelLearner(
+    time_delta=args.delta,
+    dataset_path=args.out,
+)
+
+if args.merge:
+    learner.merge_datasets(args.merge)
+    print(f"Merged existing dataset from {args.merge} ({len(learner.dataset)} samples)")
+
+print(f"Collecting {args.steps} samples (Δt={args.delta}s each) → {args.out}")
+print("Play the game — vary rod positions, pump speeds, and operating states.")
+print("Press Ctrl-C to stop early; data collected so far will be saved.")
+
+try:
+    learner.collect_data(num_steps=args.steps)
+except KeyboardInterrupt:
+    print("\nInterrupted — saving collected data...")
+
+learner.save_dataset(args.out)
+print(f"Saved {len(learner.dataset)} samples to {args.out}")

scripts/reactor_control.py

@@ -0,0 +1,739 @@
+"""Classical PID-based reactor controller with curses TUI.
+
+Architecture:
+  Core control (shared):
+    - Rod PID: keeps CORE_TEMP at setpoint via ROD_BANK_POS_0_ORDERED
+
+  Per-train control (trains 1/2/3, 0-indexed as 0/1/2 in param names):
+    - Primary pump: not touched; warns in TUI if far from suggested 65%
+    - MSCV PI: drives train power output, gated on steam availability
+    - Secondary pump feedforward: half of steam outlet + level PID
+    - Bypass: hold at 0
+
+  Auxiliary:
+    - Vacuum pump: on continuously; turned off only during retention tank drain
+    - Condenser circulation pump: fixed 25% (prevents overcooling of return water)
+    - Retention tank: drain via ejector return valve when > 75%, stop at 50%
+    - Condenser fill: run FREIGHT_PUMP_CONDENSER below 45%, stop at 60%
+
+Usage:
+    python3.14 scripts/reactor_control.py --trains 3 --target 50000
+    python3.14 scripts/reactor_control.py --trains 1 3 --target 30000 40000
+    python3.14 scripts/reactor_control.py --trains 1 2 3 --target 20000 20000 20000
+
+TUI keys:
+    0           Select core  (then +/- adjusts temp setpoint ±5°C)
+    1 / 2 / 3   Select train (then +/- adjusts target power ±5 MW; + adds if absent)
+    d           Remove selected train from control
+    g           Toggle grid-demand following
+    q / Esc     Quit
+"""
+import argparse
+import curses
+import time
+import numpy as np
+from enum import Enum
+from nucon import Nucon
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--trains',        type=int, nargs='+', default=[3])
+parser.add_argument('--target',        type=float, nargs='+', default=[50_000])
+parser.add_argument('--temp-setpoint', type=float, default=330.0)
+parser.add_argument('--dt',            type=float, default=5.0)
+parser.add_argument('--grid-follow',   action='store_true',
+                    help='Auto-set train targets from grid demand')
+parser.add_argument('--grid-buffer',   type=float, default=10.0,
+                    help='Extra MW above grid demand when grid-following (default: 5)')
+args = parser.parse_args()
+
+if len(args.target) == 1:
+    targets = {t: args.target[0] for t in args.trains}
+else:
+    if len(args.target) != len(args.trains):
+        raise ValueError("--target must have 1 value or one per --trains entry")
+    targets = dict(zip(args.trains, args.target))
+
+nucon = Nucon()
+
+# ---------------------------------------------------------------------------
+# PID controller
+# ---------------------------------------------------------------------------
+class PID:
+    def __init__(self, kp, ki, kd, out_min, out_max, integral_max=None):
+        self.kp, self.ki, self.kd = kp, ki, kd
+        self.out_min, self.out_max = out_min, out_max
+        self.integral_max = integral_max or (out_max - out_min)
+        self._integral = 0.0
+        self._prev_error = None
+
+    def step(self, error, dt):
+        self._integral = np.clip(self._integral + error * dt,
+                                 -self.integral_max, self.integral_max)
+        derivative = 0.0 if self._prev_error is None else (error - self._prev_error) / dt
+        self._prev_error = error
+        return float(np.clip(
+            self.kp * error + self.ki * self._integral + self.kd * derivative,
+            self.out_min, self.out_max))
+
+    def reset(self):
+        self._integral = 0.0
+        self._prev_error = None
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+_all_readable = None
+def _get_all_readable():
+    global _all_readable
+    if _all_readable is None:
+        _all_readable = nucon.get_all_readable()
+    return _all_readable
+
+def set_param(param_id, value):
+    param = nucon._parameters[param_id]
+    v = float(np.clip(value, param.min_val or 0, param.max_val or 100))
+    nucon.set(param, v)
+    return v
+
+def read_state(param_ids):
+    all_r = _get_all_readable()
+    raw = nucon._batch_query([p for p in param_ids if p in all_r])
+    state = {}
+    for p in param_ids:
+        if p not in all_r:
+            state[p] = 0.0
+            continue
+        try:
+            v = nucon._parse_value(all_r[p], raw.get(p, '0'))
+            state[p] = float(v.value if isinstance(v, Enum) else v)
+        except Exception:
+            state[p] = 0.0
+    return state
+
+# ---------------------------------------------------------------------------
+# Per-train controller
+# ---------------------------------------------------------------------------
+class TrainController:
+    """Controls one train (steam gen N + turbine N + generator N)."""
+
+    def __init__(self, train_num, target_kw):
+        self.n = train_num
+        self.i = train_num - 1
+        self.target_kw = target_kw
+
+        self.mscv_pid = PID(kp=0.00002, ki=0.000002, kd=0.0,
+                            out_min=-0.3, out_max=0.2, integral_max=3.0)
+        self._prev_steam_out = None
+        self.sec_pid = PID(kp=0.0005, ki=0.00005, kd=0.001,
+                           out_min=-2.0, out_max=2.0, integral_max=3.0)
+        self.sec_level_target = 25_000.0
+
+        self.prim_pump = float(nucon.get(f'COOLANT_CORE_CIRCULATION_PUMP_{self.i}_ORDERED_SPEED') or 50.0)
+        self.PRIM_PUMP_SUGGESTED = 65.0   # warn in TUI if far from this
+        self.mscv      = 9.0
+        self.sec_pump  = 40.0
+
+        set_param(f'STEAM_TURBINE_{self.i}_BYPASS_ORDERED', 0.0)
+
+        self._params = [
+            f'STEAM_GEN_{self.i}_OUTLET',
+            f'MSCV_{self.i}_OPENING_ACTUAL',
+            f'STEAM_TURBINE_{self.i}_RPM',
+            f'STEAM_TURBINE_{self.i}_BYPASS_ACTUAL',
+            f'GENERATOR_{self.i}_KW',
+            f'COOLANT_CORE_CIRCULATION_PUMP_{self.i}_ORDERED_SPEED',
+            f'COOLANT_SEC_CIRCULATION_PUMP_{self.i}_ORDERED_SPEED',
+            f'COOLANT_SEC_{self.i}_LIQUID_VOLUME',
+        ]
+
+    def params(self):
+        return self._params
+
+    def step(self, s, dt):
+        steam_out   = s[f'STEAM_GEN_{self.i}_OUTLET']
+        power_kw    = s[f'GENERATOR_{self.i}_KW']
+        power_error = self.target_kw - power_kw
+
+        # Dead-band: don't adjust MSCV when within 3% of target (avoid hunting)
+        if abs(power_error) < 0.03 * self.target_kw:
+            mscv_delta = 0.0
+            self.mscv_pid.reset()
+        else:
+            mscv_delta = self.mscv_pid.step(power_error, dt)
+        steam_rose = (self._prev_steam_out is None or
+                      steam_out >= self._prev_steam_out - 1.0)
+        if mscv_delta > 0 and not steam_rose:
+            mscv_delta = 0.0
+        self._prev_steam_out = steam_out
+        # Cap only prevents opening further — don't force MSCV down as steam fluctuates.
+        mscv_max  = max(steam_out / 8.0, 1.0)
+        new_mscv  = self.mscv + mscv_delta
+        if mscv_delta > 0:
+            new_mscv = min(new_mscv, mscv_max)
+        self.mscv = float(np.clip(new_mscv, 0.5, 100.0))
+        set_param(f'MSCV_{self.i}_OPENING_ORDERED', self.mscv)
+
+        self.prim_pump = s.get(f'COOLANT_CORE_CIRCULATION_PUMP_{self.i}_ORDERED_SPEED', self.prim_pump)
+
+        sec_ff      = steam_out / 2.0
+        level       = s[f'COOLANT_SEC_{self.i}_LIQUID_VOLUME']
+        level_error = self.sec_level_target - level
+        sec_corr    = self.sec_pid.step(level_error, dt)
+        sec_target  = float(np.clip(sec_ff + sec_corr, 5.0, 100.0))
+        self.sec_pump += 0.3 * (sec_target - self.sec_pump)
+        set_param(f'COOLANT_SEC_CIRCULATION_PUMP_{self.i}_ORDERED_SPEED', self.sec_pump)
+
+        if s[f'STEAM_TURBINE_{self.i}_BYPASS_ACTUAL'] > 1.0:
+            set_param(f'STEAM_TURBINE_{self.i}_BYPASS_ORDERED', 0.0)
+
+        return power_kw, power_error, steam_out, level, level_error
+
+# ---------------------------------------------------------------------------
+# Global controller state
+# ---------------------------------------------------------------------------
+TEMP_MAX     = 410.0
+ROD_INTERVAL = 6
+ROD_TIERS = [
+    (3.0,  0.1),
+    (8.0,  0.4),
+    (15.0, 0.8),
+    (float('inf'), 1.2),
+]
+rod_pos      = float(nucon.get('ROD_BANK_POS_0_ACTUAL') or 85.0)
+rod_cycle    = 0
+rod_integral = 0.0
+
+train_controllers = {t: TrainController(t, targets[t]) for t in args.trains}
+
+core_params = [
+    'CORE_TEMP', 'ROD_BANK_POS_0_ACTUAL',
+    'CORE_STATE_CRITICALITY',
+    'VACUUM_RETENTION_TANK_VOLUME',
+    'CONDENSER_VOLUME', 'CONDENSER_VAPOR_VOLUME',
+    'CONDENSER_VACUUM',                            # vacuum level % — monitor for pump health
+    'POWER_DEMAND_MW',
+    'CORE_PRIMARY_CIRCUIT_COOLING_TANK_VOLUME',   # pressurizer water volume
+    'COOLANT_CORE_PRIMARY_LOOP_LEVEL',             # overall primary loop fill %
+    'FREIGHT_PUMP_FEEDWATER_ACTIVE',
+]
+
+RETENTION_MAX = 40_000.0
+RETENTION_HI  = 0.75 * RETENTION_MAX
+RETENTION_MID = 0.50 * RETENTION_MAX
+
+# ---------------------------------------------------------------------------
+# Pressurizer / primary circuit constants
+# ---------------------------------------------------------------------------
+PRSR_VALVE    = 'Valvula_Pressurizer_Spray'
+# CORE_PRIMARY_CIRCUIT_COOLING_TANK_VOLUME is the pressurizer water volume.
+# Observed: 106030 = 60% → max ≈ 176717
+PRSR_VOL_MAX  = 176_717.0
+PRSR_LEVEL_LO    = 50.0   # % — open spray valve below this
+PRSR_LEVEL_CLOSE = 60.0   # % — close spray valve once level recovers
+PRSR_LEVEL_HI    = 70.0   # % — op range high (informational)
+PRIM_FILL_LO     = 80.0   # % — start feedwater pump below this (uses COOLANT_CORE_PRIMARY_LOOP_LEVEL)
+PRIM_FILL_HI     = 90.0   # % — stop feedwater pump above this
+
+# Initialise aux state from live game values so restarts are seamless.
+_init = read_state([
+    'VACUUM_RETENTION_TANK_VOLUME',
+    'STEAM_EJECTOR_CONDENSER_RETURN_VALVE_ACTUAL',
+    'CONDENSER_VOLUME', 'CONDENSER_VAPOR_VOLUME',
+    'FREIGHT_PUMP_CONDENSER_ACTIVE',
+    'CONDENSER_VACUUM_PUMP_ACTIVE',
+    'CONDENSER_CIRCULATION_PUMP_ACTIVE',
+])
+_ret_vol_init   = _init.get('VACUUM_RETENTION_TANK_VOLUME', 0.0)
+_ret_valve_init = _init.get('STEAM_EJECTOR_CONDENSER_RETURN_VALVE_ACTUAL', 0.0)
+ret_valve       = _ret_valve_init
+ret_draining    = (_ret_valve_init > 0.5 and _ret_vol_init > RETENTION_MID)
+if _ret_valve_init > 0.5 and not ret_draining:
+    set_param('STEAM_EJECTOR_CONDENSER_RETURN_VALVE', 0.0)
+    ret_valve = 0.0
+ret_prev_vol    = _ret_vol_init
+
+_cond_vol_init  = _init.get('CONDENSER_VOLUME', 0.0)
+_cond_vap_init  = _init.get('CONDENSER_VAPOR_VOLUME', 0.0)
+_cond_tot_init  = _cond_vol_init + _cond_vap_init
+_cond_pct_init  = (_cond_vol_init / _cond_tot_init * 100.0) if _cond_tot_init > 0 else 0.0
+_cond_pump_init = bool(_init.get('FREIGHT_PUMP_CONDENSER_ACTIVE', False))
+if _cond_pump_init and _cond_pct_init >= 60.0:
+    nucon.set(nucon._parameters['FREIGHT_PUMP_CONDENSER_SWITCH'], False)
+    cond_pump_on = False
+elif not _cond_pump_init and _cond_pct_init < 45.0:
+    nucon.set(nucon._parameters['FREIGHT_PUMP_CONDENSER_SWITCH'], True)
+    cond_pump_on = True
+else:
+    cond_pump_on = _cond_pump_init
+
+# Vacuum pump — keep on continuously; turn off only during retention tank drain.
+# (Opening the return valve breaks the suction path so the pump has no effect.)
+vac_pump_on = bool(_init.get('CONDENSER_VACUUM_PUMP_ACTIVE', False))
+if not vac_pump_on:
+    nucon.set(nucon._parameters['CONDENSER_VACUUM_PUMP_START_STOP'], True)
+    vac_pump_on = True
+
+# Condenser circulation pump — run at moderate speed to prevent overcooling
+# (manual §Stabilization: "prevent excessive cooling of the coolant returning to the evaporator").
+_cond_circ_on = bool(_init.get('CONDENSER_CIRCULATION_PUMP_ACTIVE', False))
+if not _cond_circ_on:
+    nucon.set(nucon._parameters['CONDENSER_CIRCULATION_PUMP_SWITCH'], True)
+    set_param('CONDENSER_CIRCULATION_PUMP_ORDERED_SPEED', 25.0)
+
+# Pressurizer spray valve — init from live state
+_prsr_live  = read_state(['CORE_PRIMARY_CIRCUIT_COOLING_TANK_VOLUME', 'COOLANT_CORE_PRIMARY_LOOP_LEVEL', 'FREIGHT_PUMP_FEEDWATER_ACTIVE'])
+_prsr_level = _prsr_live.get('CORE_PRIMARY_CIRCUIT_COOLING_TANK_VOLUME', PRSR_VOL_MAX * 0.6) / PRSR_VOL_MAX * 100.0
+_prsr_valve = nucon.get_valve(PRSR_VALVE)
+_prsr_open  = _prsr_valve.get('IsOpened', False) or _prsr_valve.get('Value', 0) > 50
+prsr_spraying = _prsr_open and _prsr_level < PRSR_LEVEL_CLOSE
+if _prsr_open and not prsr_spraying:
+    nucon.close_valve(PRSR_VALVE)
+
+feedwater_on = bool(_prsr_live.get('FREIGHT_PUMP_FEEDWATER_ACTIVE', False))
+
+# ---------------------------------------------------------------------------
+# TUI helpers
+# ---------------------------------------------------------------------------
+def _bar(pct, width=18):
+    pct = max(0.0, min(100.0, pct))
+    filled = int(pct / 100.0 * width)
+    return '█' * filled + '░' * (width - filled)
+
+def _safe_addstr(scr, row, col, text, attr=0):
+    H, W = scr.getmaxyx()
+    if row < 0 or row >= H:
+        return
+    if col < 0:
+        text = text[-col:]
+        col = 0
+    if col >= W:
+        return
+    text = text[:W - col]
+    try:
+        scr.addstr(row, col, text, attr)
+    except curses.error:
+        pass
+
+def _hline(scr, row, char='─'):
+    H, W = scr.getmaxyx()
+    if 0 <= row < H:
+        _safe_addstr(scr, row, 0, char * (W - 1))
+
+# ---------------------------------------------------------------------------
+# Main TUI loop
+# ---------------------------------------------------------------------------
+def run_controller(stdscr):
+    global rod_pos, rod_cycle, rod_integral
+    global ret_valve, ret_draining, ret_prev_vol, cond_pump_on
+    global prsr_spraying, feedwater_on
+    global vac_pump_on
+    global train_controllers, targets
+
+    curses.curs_set(0)
+    stdscr.nodelay(True)
+    curses.start_color()
+    curses.use_default_colors()
+    curses.init_pair(1, curses.COLOR_GREEN,  -1)   # good / normal
+    curses.init_pair(2, curses.COLOR_YELLOW, -1)   # warning
+    curses.init_pair(3, curses.COLOR_RED,    -1)   # alarm
+    curses.init_pair(4, curses.COLOR_CYAN,   -1)   # selected
+    curses.init_pair(5, curses.COLOR_WHITE,  curses.COLOR_BLUE)  # title bar
+
+    GREEN  = curses.color_pair(1)
+    YELLOW = curses.color_pair(2)
+    RED    = curses.color_pair(3)
+    CYAN   = curses.color_pair(4)
+    TITLE  = curses.color_pair(5)
+    BOLD   = curses.A_BOLD
+    REV    = curses.A_REVERSE
+
+    SELECTION_ORDER = [0, 1, 2, 3, 4]   # 0=core  1/2/3=trains  4=grid
+    selected_train = args.trains[0] if args.trains else 1
+    temp_setpoint  = args.temp_setpoint   # mutable; adjustable from TUI
+    temp_auto      = True                 # auto-adjust setpoint to meet total power demand
+    grid_follow    = args.grid_follow
+    # Per-train power caps: manual max each train should carry (used for proportional distribution)
+    grid_caps = {t: tc.target_kw for t, tc in train_controllers.items()}
+    cycle      = 0
+    train_data = {}
+    # display state — updated each control cycle, read by draw() at any time
+    disp = dict(s={}, dynamic_setpoint=temp_setpoint, temp_auto=temp_auto, criticality=0.0,
+                ret_pct=0.0, ret_draining=False, ret_valve=0.0,
+                cond_pct=0.0, cond_pump_on=False,
+                vac_pump_on=vac_pump_on,
+                prsr_level=_prsr_level, prsr_spraying=prsr_spraying,
+                prim_level=_prsr_live.get('COOLANT_CORE_PRIMARY_LOOP_LEVEL', 100.0),
+                feedwater_on=feedwater_on,
+                grid_follow=grid_follow, grid_demand_kw=0.0)
+
+    def rebuild_all_params():
+        p = list(core_params)
+        for tc in train_controllers.values():
+            p += tc.params()
+        return p
+
+    all_params = rebuild_all_params()
+
+    def handle_key(key):
+        nonlocal selected_train, all_params, temp_setpoint, temp_auto, grid_follow, grid_caps
+        if key in (ord('q'), 27):
+            return True   # signal quit
+        # Direct selection by number
+        elif key in (ord('0'), ord('1'), ord('2'), ord('3')):
+            selected_train = key - ord('0')
+        elif key == ord('g'):
+            selected_train = 4
+        # Up/Down cycle through selections
+        elif key == curses.KEY_UP:
+            idx = SELECTION_ORDER.index(selected_train) if selected_train in SELECTION_ORDER else 0
+            selected_train = SELECTION_ORDER[(idx - 1) % len(SELECTION_ORDER)]
+        elif key == curses.KEY_DOWN:
+            idx = SELECTION_ORDER.index(selected_train) if selected_train in SELECTION_ORDER else 0
+            selected_train = SELECTION_ORDER[(idx + 1) % len(SELECTION_ORDER)]
+        # Right/+ increase  Left/- decrease
+        elif key in (ord('+'), ord('='), curses.KEY_RIGHT):
+            if selected_train == 0 and not temp_auto:
+                temp_setpoint = min(round(temp_setpoint / 5.0) * 5.0 + 5.0, 375.0)
+            elif selected_train == 4:
+                args.grid_buffer = min(args.grid_buffer + 1.0, 100.0)
+            elif selected_train in train_controllers:
+                if grid_follow:
+                    cur = grid_caps.get(selected_train, train_controllers[selected_train].target_kw)
+                    grid_caps[selected_train] = min(round(cur / 5_000) * 5_000 + 5_000, 100_000)  # snap+step
+                else:
+                    tc = train_controllers[selected_train]
+                    tc.target_kw = min(round(tc.target_kw / 5_000) * 5_000 + 5_000, 100_000)
+                    targets[selected_train] = tc.target_kw
+                    grid_caps[selected_train] = tc.target_kw
+            elif selected_train in (1, 2, 3):
+                targets[selected_train] = 5_000
+                train_controllers[selected_train] = TrainController(selected_train, 5_000)
+                grid_caps[selected_train] = 5_000
+                all_params = rebuild_all_params()
+        elif key in (ord('-'), curses.KEY_LEFT):
+            if selected_train == 0 and not temp_auto:
+                temp_setpoint = max(round(temp_setpoint / 5.0) * 5.0 - 5.0, 250.0)
+            elif selected_train == 4:
+                args.grid_buffer = max(args.grid_buffer - 1.0, 0.0)
+            elif selected_train in train_controllers:
+                if grid_follow:
+                    cur = grid_caps.get(selected_train, train_controllers[selected_train].target_kw)
+                    grid_caps[selected_train] = max(round(cur / 5_000) * 5_000 - 5_000, 0)
+                else:
+                    tc = train_controllers[selected_train]
+                    tc.target_kw = max(round(tc.target_kw / 5_000) * 5_000 - 5_000, 0)
+                    targets[selected_train] = tc.target_kw
+                    grid_caps[selected_train] = tc.target_kw
+        elif key == ord('d'):
+            if selected_train == 0:
+                temp_auto = not temp_auto
+                disp['temp_auto'] = temp_auto
+            elif selected_train == 4:
+                grid_follow = not grid_follow
+                disp['grid_follow'] = grid_follow
+            elif selected_train in train_controllers:
+                del train_controllers[selected_train]
+                grid_caps.pop(selected_train, None)
+                if selected_train in targets:
+                    del targets[selected_train]
+                all_params = rebuild_all_params()
+                selected_train = 0 if not train_controllers else list(train_controllers.keys())[0]
+        return False
+
+    def draw():
+        s              = disp['s']
+        dynamic_setpoint = disp['dynamic_setpoint']
+        criticality    = disp['criticality']
+        ret_pct        = disp['ret_pct']
+        ret_draining   = disp['ret_draining']
+        ret_valve      = disp['ret_valve']
+        cond_pct       = disp['cond_pct']
+        cond_pump_on   = disp['cond_pump_on']
+        if not s:
+            return
+        stdscr.erase()
+        H, W = stdscr.getmaxyx()
+        row = 0
+
+        title = f" NUCLEARES CONTROLLER  ─  Cycle {cycle:5d}  ─  dt={args.dt:.0f}s "
+        _safe_addstr(stdscr, row, 0, title.ljust(W - 1), TITLE | BOLD)
+        row += 1
+
+        _hline(stdscr, row); row += 1
+        core_sel   = (selected_train == 0)
+        core_attr  = CYAN | BOLD if core_sel else BOLD
+        temp_auto_ = disp['temp_auto']
+        core_temp  = s.get('CORE_TEMP', 0.0)
+        temp_color = RED if core_temp > 370 else YELLOW if core_temp > 355 else GREEN
+        scram_str  = '  !! SCRAM !!' if core_temp > TEMP_MAX else ''
+        auto_str   = 'AUTO' if temp_auto_ else 'MAN '
+        auto_color = GREEN if temp_auto_ else YELLOW
+        _safe_addstr(stdscr, row, 2, '◆ CORE' + (' ◀' if core_sel else ''), core_attr)
+        _safe_addstr(stdscr, row, 10, f'[{auto_str}]', auto_color | BOLD)
+        _safe_addstr(stdscr, row, 16, 'Temp: ', BOLD)
+        _safe_addstr(stdscr, row, 22, f'{core_temp:6.1f}°C', temp_color | BOLD)
+        sp_color = RED if dynamic_setpoint < 306 or dynamic_setpoint > 375 else 0
+        _safe_addstr(stdscr, row, 32, f'sp=', 0)
+        _safe_addstr(stdscr, row, 35, f'{dynamic_setpoint:.0f}°C', sp_color | BOLD)
+        _safe_addstr(stdscr, row, 40,
+            f'   Rod: {s.get("ROD_BANK_POS_0_ACTUAL", 0):5.1f}   '
+            f'Crit: {criticality:+.3f}{scram_str}')
+        row += 1
+
+        for t in (1, 2, 3):
+            _hline(stdscr, row); row += 1
+            is_sel    = (t == selected_train)
+            is_active = (t in train_controllers)
+            tc        = train_controllers.get(t)
+            sel_attr  = CYAN | BOLD if is_sel else 0
+            label     = f'◆ TRAIN {t}' + (' ◀' if is_sel else '')
+            _safe_addstr(stdscr, row, 2, label, sel_attr | BOLD)
+            if is_active and t in train_data:
+                power_kw, power_error, steam_out, level, level_error = train_data[t]
+                pwr_pct   = power_kw / tc.target_kw * 100.0 if tc.target_kw > 0 else 0.0
+                pwr_color = GREEN if abs(power_error) < 2000 else YELLOW if abs(power_error) < 8000 else RED
+                cap       = grid_caps.get(t, tc.target_kw)
+                gf        = disp['grid_follow']
+                tgt_str   = (f'tgt={tc.target_kw/1000:.1f}/{cap/1000:.0f}MW'
+                             if gf and abs(tc.target_kw - cap) > 500
+                             else f'tgt={tc.target_kw/1000:.0f}MW')
+                _safe_addstr(stdscr, row, 16, 'Power: ', BOLD)
+                _safe_addstr(stdscr, row, 23, f'{power_kw/1000:5.1f} MW', pwr_color | BOLD)
+                _safe_addstr(stdscr, row, 32,
+                    f'[{_bar(pwr_pct, 14)}] {power_error/1000:+5.1f}MW  {tgt_str}')
+                row += 1
+                prim_warn  = abs(tc.prim_pump - tc.PRIM_PUMP_SUGGESTED) > 10
+                prim_attr  = YELLOW if prim_warn else 0
+                prim_str   = f'{tc.prim_pump:3.0f}%{"!" if prim_warn else " "}'
+                _safe_addstr(stdscr, row, 16,
+                    f'Steam: {steam_out:5.1f}  MSCV: {tc.mscv:4.1f}  Prim: ')
+                _safe_addstr(stdscr, row, 51, prim_str, prim_attr)
+                _safe_addstr(stdscr, row, 56,
+                    f' Sec: {tc.sec_pump:3.0f}%  Lvl: {level:.0f} (Δ{level_error:+.0f})')
+            elif not is_active:
+                hint = ' (+/Up to add)' if is_sel else ''
+                _safe_addstr(stdscr, row, 16, f'not controlled{hint}',
+                             YELLOW if is_sel else 0)
+            row += 1
+
+        _hline(stdscr, row); row += 1
+        gf         = disp['grid_follow']
+        gdkw       = disp['grid_demand_kw']
+        total_cap  = sum(grid_caps.get(t, tc.target_kw) for t, tc in train_controllers.items())
+        grid_sel   = (selected_train == 4)
+        grid_attr  = CYAN | BOLD if grid_sel else BOLD
+        gf_color   = GREEN | BOLD if gf else (CYAN | BOLD if grid_sel else 0)
+        _safe_addstr(stdscr, row, 2, '◆ GRID' + (' ◀' if grid_sel else ''), grid_attr)
+        _safe_addstr(stdscr, row, 16, f'Demand: {gdkw/1000:5.1f} MW', BOLD)
+        if gf:
+            target_total = gdkw + args.grid_buffer * 1000.0
+            _safe_addstr(stdscr, row, 34,
+                f'  AUTO  buf={args.grid_buffer:.0f}MW  '
+                f'→{target_total/1000:.1f}/{total_cap/1000:.0f}MW total', gf_color)
+        else:
+            _safe_addstr(stdscr, row, 34,
+                f'  off   buf={args.grid_buffer:.0f}MW  cap={total_cap/1000:.0f}MW', gf_color)
+        row += 1
+
+        _hline(stdscr, row); row += 1
+        ret_color  = RED if ret_pct > 75 else YELLOW if ret_pct > 60 else GREEN
+        _safe_addstr(stdscr, row, 2, '◆ RETENTION TANK  ', BOLD)
+        _safe_addstr(stdscr, row, 20, f'[{_bar(ret_pct, 20)}]', ret_color)
+        _safe_addstr(stdscr, row, 43, f' {ret_pct:4.0f}%')
+        _safe_addstr(stdscr, row, 49,
+            f'  DRAINING  valve={ret_valve:.0f}%' if ret_draining else '  OK',
+            YELLOW if ret_draining else GREEN)
+        row += 1
+        cond_vac_  = s.get('CONDENSER_VACUUM', 0.0)
+        cond_color = RED if cond_pct < 25 else YELLOW if cond_pct < 40 else GREEN
+        vac_on_    = disp.get('vac_pump_on', True)
+        vac_color  = (RED if cond_vac_ < 50 else YELLOW if cond_vac_ < 80 else GREEN) if vac_on_ else YELLOW
+        _safe_addstr(stdscr, row, 2, '◆ CONDENSER FILL  ', BOLD)
+        _safe_addstr(stdscr, row, 20, f'[{_bar(cond_pct, 20)}]', cond_color)
+        _safe_addstr(stdscr, row, 43, f' {cond_pct:4.0f}%')
+        _safe_addstr(stdscr, row, 49, '  PUMP ON' if cond_pump_on else '  OK',
+                     YELLOW if cond_pump_on else GREEN)
+        _safe_addstr(stdscr, row, 60,
+                     f'  VAC:{"OFF" if not vac_on_ else f"{cond_vac_:.0f}%"}',
+                     vac_color)
+        row += 1
+        prsr_level_  = disp['prsr_level']
+        prsr_spray_  = disp['prsr_spraying']
+        feedwater_   = disp['feedwater_on']
+        prsr_color   = RED if prsr_level_ < 40 or prsr_level_ > 80 else YELLOW if prsr_level_ < PRSR_LEVEL_LO or prsr_level_ > PRSR_LEVEL_HI else GREEN
+        _safe_addstr(stdscr, row, 2, '◆ PRESSURIZER     ', BOLD)
+        _safe_addstr(stdscr, row, 20, f'[{_bar(prsr_level_, 20)}]', prsr_color)
+        _safe_addstr(stdscr, row, 43, f' {prsr_level_:4.1f}%')
+        _safe_addstr(stdscr, row, 49, '  SPRAY ON' if prsr_spray_ else '  OK',
+                     YELLOW if prsr_spray_ else GREEN)
+        row += 1
+        prim_level_  = disp.get('prim_level', 100.0)
+        prim_color   = RED if prim_level_ < 70 else YELLOW if prim_level_ < PRIM_FILL_LO else GREEN
+        _safe_addstr(stdscr, row, 2, '◆ PRIMARY VESSEL  ', BOLD)
+        _safe_addstr(stdscr, row, 20, f'[{_bar(prim_level_, 20)}]', prim_color)
+        _safe_addstr(stdscr, row, 43, f' {prim_level_:4.1f}%')
+        _safe_addstr(stdscr, row, 49, '  FW PUMP ON' if feedwater_ else '  OK',
+                     YELLOW if feedwater_ else GREEN)
+        row += 1
+
+        if selected_train == 0:
+            adj_hint = f'←/→ sp {disp["dynamic_setpoint"]:.0f}°C±5' if not disp['temp_auto'] else f'sp={disp["dynamic_setpoint"]:.0f}°C (auto)'
+            d_hint   = f'  [d] auto {"OFF" if disp["temp_auto"] else "ON"}'
+        elif selected_train == 4:
+            adj_hint = f'←/→ buf {args.grid_buffer:.0f}MW±1'
+            d_hint   = '  [d] toggle auto'
+        elif disp['grid_follow']:
+            cap = grid_caps.get(selected_train, 0)
+            adj_hint = f'←/→ max {cap/1000:.0f}MW±5'
+            d_hint   = '  [d] remove'
+        else:
+            adj_hint = '←/→ target ±5MW'
+            d_hint   = '  [d] remove'
+        _safe_addstr(stdscr, H - 1, 0,
+            f' [↑↓] select  [0-3/g] jump  {adj_hint}{d_hint}  [q] quit '.ljust(W - 1),
+            REV)
+        stdscr.refresh()
+
+    while True:
+        t0 = time.time()
+        s = read_state(all_params)
+        cycle += 1
+        # ---- Rod control ----
+        temp_error  = s['CORE_TEMP'] - temp_setpoint
+        criticality = s.get('CORE_STATE_CRITICALITY', 0.0)
+        rod_cycle  += 1
+        if s['CORE_TEMP'] > TEMP_MAX:
+            rod_pos = 100.0
+            for tc in train_controllers.values():
+                tc.prim_pump = 90.0
+                set_param(f'COOLANT_CORE_CIRCULATION_PUMP_{tc.i}_ORDERED_SPEED', 90.0)
+        else:
+            urgent = temp_error > 5.0 or criticality > 0.3
+            if urgent or rod_cycle >= ROD_INTERVAL:
+                if rod_cycle >= ROD_INTERVAL:
+                    rod_cycle = 0
+                abs_err  = abs(temp_error)
+                max_step = next(lim for thresh, lim in ROD_TIERS if abs_err <= thresh)
+                if urgent and rod_cycle != 0:
+                    max_step = min(max_step, 0.25)
+                if not urgent:
+                    rod_integral = float(np.clip(rod_integral + 0.002 * temp_error, -3.0, 3.0))
+                else:
+                    rod_integral *= 0.5
+                raw_delta = 0.04 * temp_error + 1.0 * criticality + rod_integral
+                rod_delta = float(np.clip(raw_delta, -max_step, max_step))
+                rod_pos   = float(np.clip(s['ROD_BANK_POS_0_ACTUAL'] + rod_delta, 0.0, 100.0))
+        set_param('ROD_BANK_POS_0_ORDERED', rod_pos)
+
+        # ---- Grid-demand following ----
+        grid_demand_kw = s.get('POWER_DEMAND_MW', 0.0) * 1000.0
+        if grid_follow and train_controllers:
+            total_target_kw = grid_demand_kw + args.grid_buffer * 1000.0
+            # Distribute proportionally to each train's manual cap; never exceed cap
+            total_cap = sum(grid_caps.get(t, tc.target_kw) for t, tc in train_controllers.items())
+            if total_cap > 0:
+                for t, tc in train_controllers.items():
+                    cap = grid_caps.get(t, tc.target_kw)
+                    share = total_target_kw * (cap / total_cap)
+                    tc.target_kw = float(np.clip(share, 0.0, cap))
+
+        # ---- Per-train control ----
+        for t, tc in train_controllers.items():
+            res = tc.step(s, args.dt)
+            train_data[t] = res
+
+        # ---- Auto temp setpoint ----
+        if temp_auto and train_data:
+            total_error = sum(train_data[t][1] for t in train_data)  # sum of power_errors
+            sp_delta = float(np.clip(total_error * 0.00002, -0.5, 0.5))
+            temp_setpoint = float(np.clip(temp_setpoint + sp_delta, 306.0, 375.0))
+
+        # ---- Aux: retention tank ----
+        ret_vol = s.get('VACUUM_RETENTION_TANK_VOLUME', 0.0)
+        ret_pct = ret_vol / RETENTION_MAX * 100.0
+        if ret_draining and ret_vol <= RETENTION_MID:
+            ret_draining = False
+            ret_valve = 0.0
+            set_param('STEAM_EJECTOR_CONDENSER_RETURN_VALVE', 0.0)
+            # Drain complete — restart vacuum pump
+            if not vac_pump_on:
+                nucon.set(nucon._parameters['CONDENSER_VACUUM_PUMP_START_STOP'], True)
+                vac_pump_on = True
+        elif ret_vol > RETENTION_HI:
+            if not ret_draining:
+                # Starting drain — stop vacuum pump.
+                # The ejector return valve bypasses the suction path so the pump has no effect
+                # and wastes power; turn it off for the duration of the drain.
+                nucon.set(nucon._parameters['CONDENSER_VACUUM_PUMP_START_STOP'], False)
+                vac_pump_on = False
+            ret_draining = True
+            if ret_prev_vol is not None and ret_vol >= ret_prev_vol - 50.0:
+                ret_valve = min(ret_valve + 1.0, 50.0)
+            set_param('STEAM_EJECTOR_CONDENSER_RETURN_VALVE', ret_valve)
+        elif ret_draining:
+            set_param('STEAM_EJECTOR_CONDENSER_RETURN_VALVE', ret_valve)
+        ret_prev_vol = ret_vol
+
+        # ---- Aux: condenser fill ----
+        cond_vol = s.get('CONDENSER_VOLUME', 0.0)
+        cond_vap = s.get('CONDENSER_VAPOR_VOLUME', 0.0)
+        cond_tot = cond_vol + cond_vap
+        cond_pct = (cond_vol / cond_tot * 100.0) if cond_tot > 0 else 0.0
+        if not cond_pump_on and cond_pct < 45.0:
+            cond_pump_on = True
+            nucon.set(nucon._parameters['FREIGHT_PUMP_CONDENSER_SWITCH'], True)
+        elif cond_pump_on and cond_pct >= 60.0:
+            cond_pump_on = False
+            nucon.set(nucon._parameters['FREIGHT_PUMP_CONDENSER_SWITCH'], False)
+
+        # ---- Aux: pressurizer spray valve (level 50-70%) ----
+        prsr_level = s.get('CORE_PRIMARY_CIRCUIT_COOLING_TANK_VOLUME', PRSR_VOL_MAX * 0.6) / PRSR_VOL_MAX * 100.0
+        if not prsr_spraying and prsr_level < PRSR_LEVEL_LO:
+            prsr_spraying = True
+            nucon.open_valve(PRSR_VALVE)
+        elif prsr_spraying and prsr_level >= PRSR_LEVEL_CLOSE:
+            prsr_spraying = False
+            nucon.close_valve(PRSR_VALVE)
+        elif not prsr_spraying:
+            # Valve should be at rest — power off actuator if it's reached closed position
+            _vs = nucon.get_valve(PRSR_VALVE)
+            if _vs.get('IsClosed') and _vs.get('Actuator') != 'OFF':
+                nucon.off_valve(PRSR_VALVE)
+
+        # ---- Aux: primary circuit feedwater (overall loop fill > 80%) ----
+        prim_level = s.get('COOLANT_CORE_PRIMARY_LOOP_LEVEL', 100.0)
+        if not feedwater_on and prim_level < PRIM_FILL_LO:
+            feedwater_on = True
+            nucon.set(nucon._parameters['FREIGHT_PUMP_FEEDWATER_SWITCH'], True)
+        elif feedwater_on and prim_level >= PRIM_FILL_HI:
+            feedwater_on = False
+            nucon.set(nucon._parameters['FREIGHT_PUMP_FEEDWATER_SWITCH'], False)
+
+        # ---- Update display state and redraw ----
+        disp.update(s=s, dynamic_setpoint=temp_setpoint, temp_auto=temp_auto,
+                    criticality=criticality,
+                    ret_pct=ret_pct, ret_draining=ret_draining, ret_valve=ret_valve,
+                    cond_pct=cond_pct, cond_pump_on=cond_pump_on,
+                    vac_pump_on=vac_pump_on,
+                    prsr_level=prsr_level, prsr_spraying=prsr_spraying,
+                    prim_level=prim_level, feedwater_on=feedwater_on,
+                    grid_follow=grid_follow, grid_demand_kw=grid_demand_kw)
+        draw()
+
+        # ---- Poll input + redraw at 50 ms intervals for the rest of the cycle ----
+        sim_speed = nucon.GAME_SIM_SPEED.value or 1.0
+        deadline  = t0 + args.dt / sim_speed
+        stdscr.timeout(50)
+        while time.time() < deadline:
+            key = stdscr.getch()
+            if key == -1:
+                continue
+            if handle_key(key):
+                return
+            disp['dynamic_setpoint'] = temp_setpoint
+            disp['temp_auto']        = temp_auto
+            disp['grid_follow']      = grid_follow
+            draw()
+        stdscr.timeout(-1)
+
+curses.wrapper(run_controller)

scripts/train_sac.py

@@ -0,0 +1,151 @@
+"""SAC + HER training on kNN-GP simulator.
+
+Usage:
+    python3.14 train_sac.py
+    python3.14 train_sac.py --load /tmp/sac_nucon_knn  # hot-start from previous run
+
+Requirements:
+    - NuCon game running (for parameter metadata)
+    - /tmp/reactor_knn.pkl  (kNN-GP model)
+    - /tmp/nucon_dataset.pkl  (500-sample dataset for init_states)
+"""
+import argparse
+import pickle
+import torch
+from gymnasium.wrappers import TimeLimit
+from stable_baselines3 import SAC
+from stable_baselines3.her.her_replay_buffer import HerReplayBuffer
+from stable_baselines3.common.callbacks import CheckpointCallback
+
+from nucon.sim import NuconSimulator
+from nucon.model import ReactorDynamicsModel, MixtureModel
+from nucon.rl import NuconGoalEnv, Parameterized_Objectives, Parameterized_Terminators
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--load',    default=None,                   help='Path to existing model to hot-start from')
+parser.add_argument('--steps',   type=int, default=50_000,       help='Total timesteps (default: 50000)')
+parser.add_argument('--out',     default='/tmp/sac_nucon_knn',   help='Output path for saved model')
+parser.add_argument('--model',   default='/tmp/reactor_knn.pkl', help='Dynamics model (.pkl for kNN, .pt for NN)')
+parser.add_argument('--model2',  default=None,                   help='Second dynamics model for mixture (optional)')
+parser.add_argument('--dataset', default='/tmp/nucon_dataset.pkl', help='Dataset for init states')
+args = parser.parse_args()
+
+# ---------------------------------------------------------------------------
+# Load dynamics model(s) and dataset
+# ---------------------------------------------------------------------------
+def _load_model(path):
+    if path.endswith('.pt'):
+        ckpt = torch.load(path, weights_only=False)
+        m = ReactorDynamicsModel(ckpt['input_params'], ckpt['output_params'])
+        m.load_state_dict(ckpt['state_dict'])
+        m.eval()
+        return m
+    with open(path, 'rb') as f:
+        return pickle.load(f)
+
+dynamics_model = _load_model(args.model)
+if args.model2:
+    dynamics_model = MixtureModel(dynamics_model, _load_model(args.model2))
+
+with open(args.dataset, 'rb') as f:
+    dataset = pickle.load(f)
+
+# Seed resets to in-distribution states from dataset
+init_states = [s for _, _, s, _ in dataset]
+
+# ---------------------------------------------------------------------------
+# Build sim + env
+# ---------------------------------------------------------------------------
+sim = NuconSimulator(port=8786)
+sim.set_model(dynamics_model)
+
+BATCH_SIZE        = 2048
+MAX_EPISODE_STEPS = 200
+
+GENERATORS = ['GENERATOR_0_KW', 'GENERATOR_1_KW', 'GENERATOR_2_KW']
+POWER_RANGE = {g: (0.0, 100_000.0) for g in GENERATORS}  # per-generator kW; ~100 MW upper bound
+
+# Curated obs: physically relevant features for power control (~25 dims vs ~260 full)
+OBS_PARAMS = [
+    'CORE_TEMP', 'CORE_PRESSURE', 'CORE_STATE_CRITICALITY', 'CORE_WEAR', 'CORE_INTEGRITY',
+    'ROD_BANK_POS_0_ACTUAL', 'ROD_BANK_POS_0_ORDERED',
+    'COOLANT_CORE_FLOW_SPEED', 'COOLANT_CORE_VESSEL_TEMPERATURE',
+    'COOLANT_CORE_PRESSURE', 'COOLANT_CORE_QUANTITY_IN_VESSEL',
+    'STEAM_TURBINE_0_RPM', 'STEAM_TURBINE_0_TEMPERATURE', 'STEAM_TURBINE_0_PRESSURE',
+    'STEAM_TURBINE_1_RPM', 'STEAM_TURBINE_1_TEMPERATURE', 'STEAM_TURBINE_1_PRESSURE',
+    'STEAM_TURBINE_2_RPM', 'STEAM_TURBINE_2_TEMPERATURE', 'STEAM_TURBINE_2_PRESSURE',
+    'GENERATOR_0_V', 'GENERATOR_1_V', 'GENERATOR_2_V',
+]
+
+env = NuconGoalEnv(
+    goal_params=GENERATORS,
+    goal_range=POWER_RANGE,
+    seconds_per_step=10,
+    simulator=sim,
+    obs_params=OBS_PARAMS,
+    additional_objectives=[
+        Parameterized_Objectives['uncertainty_penalty'](start=0.3),
+        Parameterized_Objectives['temp_below_linear'](max_temp=420),
+    ],
+    additional_objective_weights=[1.0, 0.01],
+    init_states=init_states,
+    delta_action_scale=0.05,
+    goal_sampling_std=0.15,  # Gaussian delta in normalised space (~180 kW typical)
+)
+
+env = TimeLimit(env, max_episode_steps=MAX_EPISODE_STEPS)
+
+# ---------------------------------------------------------------------------
+# SAC + HER
+# learning_starts = batch_size: wait for batch_size complete (short) episodes
+# before the first gradient step. As the policy learns to stay in-dist, episodes
+# will get longer and HER has more transitions to relabel.
+# ---------------------------------------------------------------------------
+if args.load:
+    print(f"Hot-starting from {args.load}")
+    model = SAC.load(args.load, env=env, device='auto',
+                     custom_objects={'learning_rate': 3e-4, 'batch_size': BATCH_SIZE,
+                                     'tau': 0.005, 'gamma': 0.98,
+                                     'train_freq': 64, 'gradient_steps': 8,
+                                     'learning_starts': MAX_EPISODE_STEPS,
+                                     'ent_coef': 0.1})
+else:
+    model = SAC(
+        'MultiInputPolicy',
+        env,
+        replay_buffer_class=HerReplayBuffer,
+        replay_buffer_kwargs={
+            'n_sampled_goal': 4,
+            'goal_selection_strategy': 'future',
+        },
+        verbose=1,
+        learning_rate=3e-4,
+        batch_size=BATCH_SIZE,
+        tau=0.005,
+        gamma=0.98,
+        train_freq=64,
+        gradient_steps=8,
+        learning_starts=BATCH_SIZE,
+        ent_coef=0.1,  # fixed; auto-tuning diverges on this many action dims
+        device='auto',
+    )
+
+checkpoint_cb = CheckpointCallback(
+    save_freq=10_000,
+    save_path=args.out + '_checkpoints/',
+    name_prefix='sac',
+)
+
+import json, os
+
+config = {'obs_params': OBS_PARAMS}
+for save_dir in [args.out + '_checkpoints/', os.path.dirname(args.out) or '.']:
+    os.makedirs(save_dir, exist_ok=True)
+    with open(os.path.join(save_dir, 'config.json'), 'w') as f:
+        json.dump(config, f)
+
+model.learn(total_timesteps=args.steps, callback=checkpoint_cb)
+model.save(args.out)
+with open(args.out + '.json', 'w') as f:
+    json.dump(config, f)
+print(f"Saved to {args.out}.zip")

test/test.py

@@ -1,73 +0,0 @@
-import pytest
-import warnings
-from nucon import Nucon, NuconConfig, PumpStatus, PumpDryStatus, PumpOverloadStatus, BreakerStatus
-
-@pytest.fixture(scope="module")
-def nucon_setup():
-    Nucon.set_dummy_mode(False)  # Assume the game is running
-    Nucon.set_base_url("http://localhost:8080/")
-    yield
-    Nucon.set_dummy_mode(True)
-
-def test_read_all_parameters(nucon_setup):
-    all_params = Nucon.get_all()
-    assert len(all_params) == len(Nucon)
-    for param, value in all_params.items():
-        assert isinstance(value, param.param_type), f"Parameter {param.name} has incorrect type. Expected {param.param_type}, got {type(value)}"
-        if param.param_type == float and value.is_integer():
-            warnings.warn(f"Parameter {param.name} is a float but has an integer value: {value}")
-
-def test_write_writable_parameters(nucon_setup):
-    writable_params = Nucon.get_all_writable()
-    for param in writable_params:
-        current_value = param.value
-        param.value = current_value
-        assert param.value == current_value, f"Failed to write to parameter {param.name}"
-
-def test_non_writable_parameters(nucon_setup):
-    non_writable_params = [param for param in Nucon if not param.is_writable]
-    for param in non_writable_params:
-        # Test that normal set raises an error
-        with pytest.raises(ValueError, match=f"Parameter {param.name} is not writable"):
-            param.value = param.value  # Attempt to write the current value
-
-        # Test that force_set is refused by the webserver
-        current_value = param.value
-        with pytest.raises(Exception, match=f"Failed to set parameter {param.name}"):
-            Nucon.set(param, current_value, force=True)
-
-def test_enum_parameters(nucon_setup):
-    pump_status = Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value
-    assert isinstance(pump_status, PumpStatus)
-
-    dry_status = Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value
-    assert isinstance(dry_status, PumpDryStatus)
-
-    overload_status = Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value
-    assert isinstance(overload_status, PumpOverloadStatus)
-
-    breaker_status = Nucon.GENERATOR_0_BREAKER.value
-    assert isinstance(breaker_status, BreakerStatus)
-
-def test_custom_truthy_values(nucon_setup):
-    assert bool(Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value) == (Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value in [PumpStatus.ACTIVE_NO_SPEED_REACHED, PumpStatus.ACTIVE_SPEED_REACHED])
-    assert bool(Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value) == (Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value == PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID)
-    assert bool(Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value) == (Nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value == PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD)
-    assert bool(Nucon.GENERATOR_0_BREAKER.value) == (Nucon.GENERATOR_0_BREAKER.value == BreakerStatus.OPEN)
-
-def test_get_multiple_parameters(nucon_setup):
-    params_to_get = [Nucon.CORE_TEMP, Nucon.CORE_PRESSURE, Nucon.RODS_POS_ACTUAL]
-    multiple_params = Nucon.get_multiple(params_to_get)
-    assert len(multiple_params) == len(params_to_get)
-    for param, value in multiple_params.items():
-        assert isinstance(value, param.param_type)
-
-def test_base_url_setting():
-    original_url = NuconConfig.base_url
-    new_url = "http://newlocalhost:9090/"
-    Nucon.set_base_url(new_url)
-    assert NuconConfig.base_url == new_url
-    Nucon.set_base_url(original_url)
-
-if __name__ == "__main__":
-    pytest.main()

test/test_core.py

@@ -0,0 +1,86 @@
+import pytest
+import warnings
+from nucon import Nucon, PumpStatus, PumpDryStatus, PumpOverloadStatus, BreakerStatus
+
+WARN_FLOAT_COULD_BE_INT = False
+
+
+@pytest.fixture(scope="function")
+def nucon():
+    """Create a fresh Nucon instance for each test"""
+    return Nucon()
+
+def test_read_all_parameters(nucon):
+    all_params = nucon.get_all()
+    assert len(all_params) == len(nucon)
+    for param, value in all_params.items():
+        param_type = nucon.get_type(param)
+        if value is None:
+            continue  # Some params return null/empty when subsystem not installed
+        assert isinstance(value, param_type), f"Parameter {param} has incorrect type. Expected {param_type}, got {type(value)}"
+        if param_type == float and value.is_integer() and WARN_FLOAT_COULD_BE_INT:
+            warnings.warn(f"Parameter {param} is a float but has an integer value: {value}")
+
+def test_write_writable_parameters(nucon):
+    writable_params = nucon.get_all_writable()
+    for param in writable_params.values():
+        if not param.is_readable:
+            continue  # Skip write-only params (can't read back to verify, and actions like SCRAM are dangerous to trigger)
+        current_value = param.value
+        if current_value is None:
+            continue  # Skip params that return null (subsystem not installed)
+        param.value = current_value
+        assert param.value == current_value, f"Failed to write to parameter {param.id}"
+
+def test_non_writable_parameters(nucon):
+    non_writable_params = [param for param in nucon.get_all_readable().values() if not param.is_writable]
+    for param in non_writable_params:
+        # Test that normal set raises an error
+        with pytest.raises(ValueError, match=f"Parameter {param.id} is not writable"):
+            param.value = param.value  # Attempt to write the current value
+
+        # Note: the game accepts force writes silently (does not return an error)
+
+def test_enum_parameters(nucon):
+    pump_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value
+    assert isinstance(pump_status, PumpStatus)
+
+    dry_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value
+    assert isinstance(dry_status, PumpDryStatus)
+
+    overload_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value
+    assert isinstance(overload_status, PumpOverloadStatus)
+
+    breaker_status = nucon.GENERATOR_0_BREAKER.value
+    assert isinstance(breaker_status, BreakerStatus)
+
+def test_custom_truthy_values(nucon):
+    assert bool(nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value) == (nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value in [PumpStatus.ACTIVE_NO_SPEED_REACHED, PumpStatus.ACTIVE_SPEED_REACHED])
+    #assert bool(nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value) == (nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value == PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID)
+    #assert bool(nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value) == (nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value == PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD)
+    assert bool(nucon.GENERATOR_0_BREAKER.value) == (nucon.GENERATOR_0_BREAKER.value == BreakerStatus.OPEN)
+
+def test_get_multiple_parameters(nucon):
+    params_to_get = [nucon.CORE_TEMP, nucon.CORE_PRESSURE, nucon.TIME]
+    multiple_params = nucon.get_multiple(params_to_get)
+    assert len(multiple_params) == len(params_to_get)
+    for param, value in multiple_params.items():
+        param_type = nucon.get_type(param)
+        assert isinstance(value, param_type)
+
+def test_param_coverage(nucon):
+    raw_game_vars = nucon.get_game_variable_names()
+    game_vars = set(v.upper() for v in raw_game_vars)
+    our_vars = set(nucon._parameters.keys())
+
+    assert game_vars, f"Got empty variable list from game. Raw sample: {raw_game_vars[:5]}"
+
+    missing_from_game = our_vars - game_vars
+    assert not missing_from_game, f"Params defined in NuCon but not found in game: {missing_from_game}\nGame var sample: {sorted(game_vars)[:10]}"
+
+    not_in_nucon = game_vars - our_vars
+    if not_in_nucon:
+        warnings.warn(f"Game exposes {len(not_in_nucon)} params not defined in NuCon: {sorted(not_in_nucon)}")
+
+if __name__ == "__main__":
+    pytest.main()

test/test_sim.py

@@ -0,0 +1,113 @@
+import pytest
+import time
+from nucon import Nucon, PumpStatus, PumpDryStatus, PumpOverloadStatus, BreakerStatus
+from nucon.sim import NuconSimulator, OperatingState
+
+@pytest.fixture(scope="module")
+def simulator_setup():
+    simulator = NuconSimulator(port=8786)
+    time.sleep(1)  # Give the simulator time to start
+    nucon = Nucon(port=8786)
+    return simulator, nucon
+
+def test_simulator_initialization(simulator_setup):
+    simulator, nucon = simulator_setup
+    assert simulator is not None
+    assert nucon is not None
+
+def test_read_all_parameters(simulator_setup):
+    _, nucon = simulator_setup
+    all_params = nucon.get_all()
+    assert len(all_params) == len(nucon)
+    for param_id, value in all_params.items():
+        param = nucon[param_id]
+        assert isinstance(value, param.param_type), f"Parameter {param.id} has incorrect type. Expected {param.param_type}, got {type(value)}"
+
+def test_write_writable_parameters(simulator_setup):
+    _, nucon = simulator_setup
+    writable_params = nucon.get_all_writable()
+    for param in writable_params.values():
+        current_value = param.value
+        if param.param_type == float:
+            new_value = current_value + 1.0
+        elif param.param_type == int:
+            new_value = current_value + 1
+        elif param.param_type == bool:
+            new_value = not current_value
+        elif issubclass(param.param_type, Enum):
+            new_value = list(param.param_type)[0]
+        else:
+            continue  # Skip if we can't determine a new value
+
+        param.value = new_value
+        assert param.value == new_value, f"Failed to write to parameter {param.id}"
+
+def test_non_writable_parameters(simulator_setup):
+    _, nucon = simulator_setup
+    non_writable_params = [param for param in nucon.get_all_readable().values() if not param.is_writable]
+    for param in non_writable_params:
+        with pytest.raises(ValueError, match=f"Parameter {param.id} is not writable"):
+            param.value = param.value  # Attempt to write the current value
+
+def test_enum_parameters(simulator_setup):
+    _, nucon = simulator_setup
+    pump_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value
+    assert isinstance(pump_status, PumpStatus)
+
+    dry_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value
+    assert isinstance(dry_status, PumpDryStatus)
+
+    overload_status = nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value
+    assert isinstance(overload_status, PumpOverloadStatus)
+
+    breaker_status = nucon.GENERATOR_0_BREAKER.value
+    assert isinstance(breaker_status, BreakerStatus)
+
+def test_custom_truthy_values(simulator_setup):
+    _, nucon = simulator_setup
+    assert bool(nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value) == (nucon.COOLANT_CORE_CIRCULATION_PUMP_0_STATUS.value in [PumpStatus.ACTIVE_NO_SPEED_REACHED, PumpStatus.ACTIVE_SPEED_REACHED])
+    assert bool(nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value) == (nucon.COOLANT_CORE_CIRCULATION_PUMP_0_DRY_STATUS.value == PumpDryStatus.INACTIVE_OR_ACTIVE_WITH_FLUID)
+    assert bool(nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value) == (nucon.COOLANT_CORE_CIRCULATION_PUMP_0_OVERLOAD_STATUS.value == PumpOverloadStatus.INACTIVE_OR_ACTIVE_NO_OVERLOAD)
+    assert bool(nucon.GENERATOR_0_BREAKER.value) == (nucon.GENERATOR_0_BREAKER.value == BreakerStatus.OPEN)
+
+def test_get_multiple_parameters(simulator_setup):
+    _, nucon = simulator_setup
+    params_to_get = [nucon.CORE_TEMP, nucon.CORE_PRESSURE, nucon.RODS_POS_ACTUAL]
+    multiple_params = nucon.get_multiple(params_to_get)
+    assert len(multiple_params) == len(params_to_get)
+    for param_id, value in multiple_params.items():
+        assert isinstance(value, nucon[param_id].param_type)
+
+def test_simulator_update(simulator_setup):
+    simulator, nucon = simulator_setup
+    initial_temp = nucon.CORE_TEMP.value
+    simulator.update(10)  # Update for 10 seconds
+    new_temp = nucon.CORE_TEMP.value
+    assert new_temp != initial_temp, "Core temperature should change after update"
+
+def test_set_operating_state(simulator_setup):
+    simulator, nucon = simulator_setup
+    simulator._set_state(OperatingState.NOMINAL)
+    time.sleep(1)  # Give the simulator time to update
+    assert nucon.CORE_STATE.value == 'NOMINAL'
+    assert 290 <= nucon.CORE_TEMP.value <= 370
+
+def test_allow_all_writes(simulator_setup):
+    simulator, nucon = simulator_setup
+    simulator.set_allow_all_writes(True)
+    non_writable_param = next(param for param in nucon.get_all_readable().values() if not param.is_writable)
+    current_value = non_writable_param.value
+    new_value = current_value + 1 if isinstance(current_value, (int, float)) else not current_value
+    non_writable_param.value = new_value
+    assert non_writable_param.value == new_value, f"Failed to write to non-writable parameter {non_writable_param.id} when allow_all_writes is True"
+
+def test_simulator_consistency(simulator_setup):
+    simulator, nucon = simulator_setup
+    for _ in range(10):
+        simulator.update(1)
+        all_params = nucon.get_all()
+        for param_id, value in all_params.items():
+            assert value == getattr(simulator.parameters, param_id), f"Inconsistency found for parameter {param_id}"
+
+if __name__ == "__main__":
+    pytest.main()