feat: SAC+HER training on kNN-GP sim with direct bypass and scripts/

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

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

nucon/model.py

@@ -98,11 +98,11 @@ class ReactorKNNModel:
         self._std  = self._raw_states.std(axis=0) + 1e-8
         self._states = (self._raw_states - self._mean) / self._std
 
-    def _lookup(self, state_dict: Dict):
-        """Return (s_norm, idx, k) for the k nearest neighbours."""
-        s      = np.array([state_dict[p] for p in self.input_params], dtype=np.float32)
+    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
-        dists  = np.linalg.norm(self._states - s_norm, axis=1)
+        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
@@ -122,22 +122,22 @@ class ReactorKNNModel:
         if self._states is None:
             raise ValueError("Model not fitted. Call fit(dataset) first.")
 
-        s_norm, idx, k = self._lookup(state_dict)
+        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 (vectorised): k(a,b) = exp(-0.5 ||a-b||^2)
-        def rbf_matrix(A, B):
-            diff = A[:, None, :] - B[None, :, :]          # (|A|, |B|, d)
-            return np.exp(-0.5 * (diff ** 2).sum(axis=-1))  # (|A|, |B|)
+        # 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_matrix(X, X) + 1e-4 * np.eye(k)      # (k, k)
-        k_star = rbf_matrix(s_norm[None, :], X)[0]         # (k,)
+        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                     # (d_out,)
+        K_inv      = np.linalg.inv(K)
+        mean_rates = k_star @ K_inv @ Y
 
-        # Posterior variance (scalar, shared across all output dims)
         var = max(0.0, 1.0 - float(k_star @ K_inv @ k_star))
         std = float(np.sqrt(var))
 

nucon/rl.py

@@ -49,13 +49,15 @@ Parameterized_Terminators = {
 # Internal helpers
 # ---------------------------------------------------------------------------
 
-def _build_flat_action_space(nucon, obs_param_set=None):
-    """Return (Box, ordered_param_ids) for all writable, readable, non-cheat params.
+def _build_flat_action_space(nucon, obs_param_set=None, delta_action_scale=None):
+    """Return (Box, ordered_param_ids, param_ranges).
 
-    If obs_param_set is provided, only include params in that set.
+    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 = [], []
+    lows, highs, ranges = [], [], []
     for param_id, param in nucon.get_all_writable().items():
         if not param.is_readable or param.is_cheat:
             continue
@@ -69,9 +71,15 @@ def _build_flat_action_space(nucon, obs_param_set=None):
         params.append(param_id)
         lows.append(sp.low[0])
         highs.append(sp.high[0])
-    box = spaces.Box(low=np.array(lows, dtype=np.float32),
-                     high=np.array(highs, dtype=np.float32), dtype=np.float32)
-    return box, params
+        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):
@@ -96,12 +104,15 @@ def _build_param_space(param):
 def _apply_action(nucon, action):
     for param_id, value in action.items():
         param = nucon._parameters[param_id]
-        if issubclass(param.param_type, Enum):
-            value = param.param_type(int(np.asarray(value).flat[0]))
+        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(np.asarray(value).flat[0])
-        if param.min_val is not None and param.max_val is not None:
-            value = np.clip(value, param.min_val, param.max_val)
+            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)
 
 
@@ -138,7 +149,8 @@ class NuconEnv(gym.Env):
                 obs_spaces[param_id] = sp
         self.observation_space = spaces.Dict(obs_spaces)
 
-        self.action_space, self._action_params = _build_flat_action_space(self.nucon)
+        self.action_space, self._action_params, self._action_lows, self._action_ranges = \
+            _build_flat_action_space(self.nucon)
 
         self.objectives = []
         self.terminators = []
@@ -272,11 +284,14 @@ class NuconGoalEnv(gym.Env):
         additional_objectives=None,
         additional_objective_weights=None,
         obs_params=None,
+        init_states=None,
+        delta_action_scale=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)
@@ -339,12 +354,14 @@ class NuconGoalEnv(gym.Env):
         })
 
         # Action space: writable params within the obs param set (flat Box for SB3 compatibility).
-        self.action_space, self._action_params = _build_flat_action_space(self.nucon, set(base_params))
+        self.action_space, self._action_params, self._action_lows, self._action_ranges = \
+            _build_flat_action_space(self.nucon, set(base_params), 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._desired_goal = np.zeros(n_goals, dtype=np.float32)
         self._total_steps = 0
 
@@ -367,20 +384,37 @@ class NuconGoalEnv(gym.Env):
     def _read_obs(self, sim_uncertainty=None):
         """Return (gym_obs_dict, reward_obs_dict).
 
-        gym_obs_dict  — flat Box observation for the policy (no SIM_UNCERTAINTY).
-        reward_obs_dict — same values plus SIM_UNCERTAINTY for objectives/terminators/reward_fn.
+        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 param_id in self._obs_params:
-            value = self.nucon.get(param_id)
-            if isinstance(value, Enum):
-                value = value.value
-            reward_obs[param_id] = float(value) if value 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)
-        achieved = self._read_goal_values()
+        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
@@ -390,11 +424,47 @@ class NuconGoalEnv(gym.Env):
         self._total_steps = 0
         rng = np.random.default_rng(seed)
         self._desired_goal = rng.uniform(0.0, 1.0, size=len(self.goal_params)).astype(np.float32)
+        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
         gym_obs, _ = self._read_obs()
         return gym_obs, {}
 
     def step(self, action):
-        _apply_action(self.nucon, _unflatten_action(action, self._action_params))
+        flat = np.asarray(action, dtype=np.float32)
+        if self._delta_action_scale is not None:
+            # Compute absolute values from deltas, reading current state directly if possible
+            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:
+                        current = 0.0  # fallback; batch read not worth it for actions alone
+                    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)

nucon/sim.py

@@ -261,14 +261,13 @@ class NuconSimulator:
             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(self.parameters, param_id, None)
+            value = getattr(params, param_id, None)
             if isinstance(value, Enum):
                 value = value.value
-            if value is None:
-                value = 0.0  # fallback for params not initialised in sim state
-            state[param_id] = value
+            state[param_id] = 0.0 if value is None else value
 
         # Forward pass
         uncertainty = None
@@ -280,12 +279,9 @@ class NuconSimulator:
         else:
             next_state = self.model.forward(state, time_step)
 
-        # Update only the output params the model predicts
+        # Write outputs directly — bypass sim.set() type-checking overhead
         for param_id, value in next_state.items():
-            try:
-                self.set(param_id, value, force=True)
-            except (ValueError, KeyError):
-                pass  # ignore params that can't be set (type mismatch, unknown)
+            setattr(params, param_id, value)
 
         return uncertainty