shaibu01/Titan-Engine
0
1"""2TITAN QUANT GATEKEEPER (v53.0 - DUAL-HEAD + 50-EXPERT + REGIME-AWARE INFERENCE)3TYPE: Runtime Inference Engine4ARCH: 50x PyTorch GRU + BatchNorm + RegimeGate + Policy Head + Value Head5FUNCTION: Regime-aware deep neural ensembling.6 Returns Consensus Probability, Epistemic Uncertainty, and Ensemble Value Score.7UPGRADES (v53.0 over v52.1):8 1. Architecture upgraded from 8D single-head to 16D dual-head (policy + value).9 2. RegimeGate added: regime-conditional feature scaling via HMM regime_signal.10 3. Expert count expanded from 5 to 50 (loads all available brain_models).11 4. Hyperparameters loaded dynamically from optimal_params.json (no hardcoding).12 5. Legacy state-dict remapping: fc2.* → policy_head.* for old checkpoints.13 6. Tensor construction now uses alpha_physics 16D feature dict (not old 8D C++ path).14 7. Device hierarchy: CUDA → MPS → CPU (cloud-portable).15 8. Returns (policy_prob, uncertainty, value_score) for PSO steering compatibility.16 9. is_trained threshold is dynamic (>= 1 expert loaded).17 10. Hot-reload fully preserved and updated for 50-expert architecture.18 11. Persistent Storage: Wires MODELS_DIR to /data to survive Space restarts.19"""20 21from __future__ import annotations22 23import json24import logging25import os26import warnings27from typing import Dict, List, Optional, Tuple28 29import numpy as np30import torch31import torch.nn as nn32 33warnings.filterwarnings("ignore")34 35logger = logging.getLogger("GATEKEEPER")36logger.setLevel(logging.INFO)37if not logger.handlers:38 sh = logging.StreamHandler()39 sh.setFormatter(logging.Formatter("%(asctime)s | GATEKEEPER | %(message)s"))40 logger.addHandler(sh)41 42# ==============================================================================43# DEVICE HIERARCHY: CUDA → MPS → CPU44# ==============================================================================45if torch.cuda.is_available():46 DEVICE = torch.device("cuda")47elif torch.backends.mps.is_available():48 DEVICE = torch.device("mps")49else:50 DEVICE = torch.device("cpu")51 torch.set_num_threads(2)52 53logger.info(f"GATEKEEPER: Running on device={DEVICE}")54 55# ==============================================================================56# CONSTANTS57# ==============================================================================58def _models_dir() -> str:59 configured = os.getenv("TITAN_MODELS_DIR", "").strip()60 if configured:61 return os.path.abspath(os.path.expanduser(configured))62 try:63 import titan_storage64 return str(titan_storage.local_path("brain_models"))65 except Exception:66 return os.path.join(os.getcwd(), ".titan_engine_storage", "brain_models")67 68 69MODELS_DIR = _models_dir()70os.makedirs(MODELS_DIR, exist_ok=True)71PARAMS_FILE = os.path.join(MODELS_DIR, "optimal_params.json")72NUM_EXPERTS_MAX = 50 # Maximum experts to attempt loading73INPUT_SIZE = 16 # 16D feature tensor from alpha_physics v51.174SEQ_LENGTH = 30 # Default sequence length (overridden by optimal_params.json)75N_REGIMES = 4 # Bull, Chop, Bear, Crash76 77# Minimum experts needed before gatekeeper considers itself operational78MIN_EXPERTS_OPERATIONAL = 179 80FEATURE_NAMES = [81 "stark_s", "stark_l", "wave", "cairo", "frama", "flux",82 "echo_s", "echo_l", "ridge", "spectra", "gap", "vix",83 "htf_trend", "vamp", "bbp", "regime_signal",84]85 86# ==============================================================================87# 1. REGIME GATE88# ==============================================================================89class RegimeGate(nn.Module):90 """91 Regime-Conditional Feature Importance Scaling.92 Learns per-regime soft on/off gates for each of the 16 input channels.93 Initialised to 1.0 so the gate is transparent before training.94 """95 def __init__(self, input_size: int = INPUT_SIZE, n_regimes: int = N_REGIMES):96 super().__init__()97 self.gate_weights = nn.Parameter(torch.ones(n_regimes, input_size))98 99 def forward(self, x: torch.Tensor, regime_idx: torch.Tensor) -> torch.Tensor:100 weights = self.gate_weights[regime_idx] # (batch, features)101 weights = torch.sigmoid(weights) # Soft gates ∈ (0, 1)102 return x * weights.unsqueeze(1) # (batch, 1, features) broadcast103 104 105# ==============================================================================106# 2. DUAL-HEAD TITANNET (matches trainer v203.3 exactly)107# ==============================================================================108class TitanNet(nn.Module):109 """110 GRU-based dual-head network.111 Track A (policy_head): Sigmoid output — probability of upward price move.112 Track B (value_head): Tanh output — expected step-by-step state quality.113 forward() defaults to policy-only for backward compatibility.114 """115 def __init__(116 self,117 input_size: int = INPUT_SIZE,118 hidden_size: int = 64,119 num_layers: int = 2,120 dropout_rate: float = 0.2,121 n_regimes: int = N_REGIMES,122 ):123 super().__init__()124 self.regime_gate = RegimeGate(input_size=input_size, n_regimes=n_regimes)125 self.gru = nn.GRU(126 input_size,127 hidden_size=hidden_size,128 num_layers=num_layers,129 batch_first=True,130 dropout=dropout_rate if num_layers > 1 else 0.0,131 )132 self.bn = nn.BatchNorm1d(hidden_size)133 self.fc1 = nn.Linear(hidden_size, hidden_size // 2)134 self.relu = nn.ReLU()135 self.drop = nn.Dropout(dropout_rate)136 137 self.policy_head = nn.Linear(hidden_size // 2, 1)138 self.value_head = nn.Linear(hidden_size // 2, 1)139 self.sig = nn.Sigmoid()140 141 def _encode(142 self,143 x: torch.Tensor,144 regime_idx: Optional[torch.Tensor] = None,145 ) -> torch.Tensor:146 if regime_idx is None:147 regime_idx = torch.zeros(x.size(0), dtype=torch.long, device=x.device)148 x = self.regime_gate(x, regime_idx)149 out, _ = self.gru(x)150 out = out[:, -1, :]151 out = self.bn(out)152 out = self.fc1(out)153 out = self.relu(out)154 out = self.drop(out)155 return out156 157 def forward(158 self,159 x: torch.Tensor,160 regime_idx: Optional[torch.Tensor] = None,161 return_value: bool = False,162 ):163 """164 Returns policy_prob (float tensor) by default.165 If return_value=True, returns (policy_prob, value_estimate).166 """167 hidden = self._encode(x, regime_idx=regime_idx)168 policy_prob = self.sig(self.policy_head(hidden))169 if not return_value:170 return policy_prob171 value_estimate = torch.tanh(self.value_head(hidden))172 return policy_prob, value_estimate173 174 175# ==============================================================================176# 3. HELPER FUNCTIONS177# ==============================================================================178def infer_regime_idx_from_sequence(x: torch.Tensor) -> torch.Tensor:179 """180 Derives discrete regime bucket from the continuous regime_signal channel181 (last feature of the last timestep, index -1 in both seq and feature dims).182 Buckets: [0.00-0.25) → 0 Bull, [0.25-0.50) → 1 Chop,183 [0.50-0.75) → 2 Bear, [0.75-1.00] → 3 Crash184 """185 regime_signal = torch.clamp(x[:, -1, -1], 0.0, 1.0)186 bins = torch.tensor([0.25, 0.50, 0.75], device=x.device)187 return torch.bucketize(regime_signal, bins).long()188 189def remap_legacy_state_dict(190 state_dict: Dict[str, torch.Tensor]191) -> Dict[str, torch.Tensor]:192 """193 Maps old single-head checkpoint keys (fc2.*) to new dual-head keys (policy_head.*).194 Safe to call on already-upgraded checkpoints — no-op if fc2 keys are absent.195 """196 new_sd = dict(state_dict)197 if "fc2.weight" in new_sd and "policy_head.weight" not in new_sd:198 new_sd["policy_head.weight"] = new_sd.pop("fc2.weight")199 if "fc2.bias" in new_sd and "policy_head.bias" not in new_sd:200 new_sd["policy_head.bias"] = new_sd.pop("fc2.bias")201 return new_sd202 203def safe_load_model_state(204 model: nn.Module,205 state_dict: Dict[str, torch.Tensor],206) -> Tuple[List[str], List[str]]:207 """208 Loads state dict with legacy remapping and strict=False so missing value_head209 keys (from old single-head checkpoints) don't crash the loader.210 Returns (missing_keys, unexpected_keys) for diagnostic logging.211 """212 remapped = remap_legacy_state_dict(state_dict)213 result = model.load_state_dict(remapped, strict=False)214 missing = list(getattr(result, "missing_keys", []))215 unexpected = list(getattr(result, "unexpected_keys", []))216 return missing, unexpected217 218def load_optimal_params() -> dict:219 """220 Reads hyperparameters from optimal_params.json written by neural_trainer v203.3.221 Falls back to safe defaults if the file is absent or malformed.222 """223 defaults = {224 "hidden_size": 64,225 "num_layers": 2,226 "dropout_rate": 0.2,227 "seq_length": SEQ_LENGTH,228 }229 if not os.path.exists(PARAMS_FILE):230 logger.warning(f"optimal_params.json not found at {PARAMS_FILE}. Using defaults.")231 return defaults232 try:233 with open(PARAMS_FILE, "r", encoding="utf-8") as f:234 data = json.load(f)235 236 params = {237 "hidden_size": int(data.get("hidden_size", defaults["hidden_size"])),238 "num_layers": int(data.get("num_layers", defaults["num_layers"])),239 "dropout_rate": float(data.get("dropout_rate", defaults["dropout_rate"])),240 "seq_length": int(data.get("seq_length", defaults["seq_length"])),241 }242 logger.info(f"GATEKEEPER: Loaded optimal_params.json → {params}")243 return params244 except Exception as e:245 logger.warning(f"optimal_params.json parse error: {e}. Using defaults.")246 return defaults247 248def build_tensor_from_physics(249 tensor_dict: dict,250 seq_length: int = SEQ_LENGTH,251) -> Optional[torch.Tensor]:252 """253 Converts a single alpha_physics tensor_dict (16 scalar features) into a254 PyTorch tensor shaped [1, seq_length, 16] by repeating the snapshot across255 the sequence dimension.256 In live inference the caller only has the LATEST bar's feature snapshot,257 not a full rolling window. Repeating across seq_length is a valid approximation258 for GRU inference — the GRU's hidden state converges quickly and the regime gate259 ensures the correct feature-importance scaling is applied regardless.260 For higher fidelity, the caller can pass a pre-built [seq_length, 16] history261 array via build_tensor_from_history().262 """263 if not tensor_dict or len(tensor_dict) < INPUT_SIZE:264 return None265 try:266 feature_vector = np.array(267 [float(tensor_dict.get(name, 0.0)) for name in FEATURE_NAMES],268 dtype=np.float32,269 )270 271 # Repeat the snapshot across seq_length timesteps: shape [seq_length, 16]272 sequence = np.tile(feature_vector, (seq_length, 1))273 274 # Add batch dimension: [1, seq_length, 16]275 tensor = torch.tensor(sequence, dtype=torch.float32).unsqueeze(0).to(DEVICE)276 return tensor277 except Exception as e:278 logger.error(f"build_tensor_from_physics error: {e}")279 return None280 281def build_tensor_from_history(282 feature_history: List[List[float]],283 seq_length: int = SEQ_LENGTH,284) -> Optional[torch.Tensor]:285 """286 Converts a rolling history of 16D feature snapshots (list of lists) into287 a [1, seq_length, 16] tensor using the last seq_length rows.288 This is the preferred path when the caller maintains a rolling feature buffer.289 """290 if not feature_history or len(feature_history[0]) != INPUT_SIZE:291 return None292 try:293 arr = np.array(feature_history[-seq_length:], dtype=np.float32)294 if arr.shape[0] < seq_length:295 # Pad at the front with the first row if history is shorter than seq_length296 pad = np.tile(arr[0], (seq_length - arr.shape[0], 1))297 arr = np.vstack([pad, arr])298 299 tensor = torch.tensor(arr, dtype=torch.float32).unsqueeze(0).to(DEVICE)300 return tensor301 except Exception as e:302 logger.error(f"build_tensor_from_history error: {e}")303 return None304 305# ==============================================================================306# 4. COUNCIL OF QUANTS — 50-EXPERT DUAL-HEAD INFERENCE ENGINE307# ==============================================================================308class CouncilOfQuants:309 """310 Loads up to 50 dual-head GRU experts trained by neural_trainer v203.3.311 Provides regime-aware ensemble inference returning:312 - consensus_prob : mean policy probability across all experts313 - uncertainty : std of policy probabilities (epistemic disagreement)314 - ensemble_value : mean value-head score (for PSO steering in titan_core)315 """316 def __init__(self):317 self.council_brains: List[nn.Module] = []318 self.models_dir = MODELS_DIR319 self.is_trained = False320 self.params = load_optimal_params()321 self.seq_length = self.params["seq_length"]322 323 self._load_council()324 325 # --------------------------------------------------------------------------326 # Internal loader327 # --------------------------------------------------------------------------328 def _load_council(self) -> None:329 """330 Scans brain_models/ for expert_*.pth files and loads all available ones.331 Applies legacy state-dict remapping for older single-head checkpoints.332 """333 self.council_brains.clear()334 self.is_trained = False335 336 if not os.path.exists(self.models_dir):337 logger.warning(f"GATEKEEPER: models_dir '{self.models_dir}' not found. Awaiting Forge.")338 return339 340 loaded = 0341 failed = 0342 343 for i in range(1, NUM_EXPERTS_MAX + 1):344 path = os.path.join(self.models_dir, f"expert_{i}_gru.pth")345 if not os.path.exists(path):346 continue347 348 try:349 model = TitanNet(350 input_size=INPUT_SIZE,351 hidden_size=self.params["hidden_size"],352 num_layers=self.params["num_layers"],353 dropout_rate=self.params["dropout_rate"],354 n_regimes=N_REGIMES355 ).to(DEVICE)356 357 state_dict = torch.load(path, map_location=DEVICE)358 missing, unexpected = safe_load_model_state(model, state_dict)359 360 # Only log verbose missing keys if they are NOT value_head (which is expected for legacy)361 important_missing = [k for k in missing if "value_head" not in k]362 if important_missing:363 logger.debug(f"Expert {i} missing critical keys: {important_missing}")364 if unexpected:365 logger.debug(f"Expert {i} unexpected keys: {unexpected}")366 367 model.eval()368 self.council_brains.append(model)369 loaded += 1370 371 except Exception as e:372 logger.error(f"Failed to load expert {i}: {e}")373 failed += 1374 375 if loaded >= MIN_EXPERTS_OPERATIONAL:376 self.is_trained = True377 logger.info(f"🧠 [GATEKEEPER] {loaded}-Expert Neural Council Loaded into RAM. (Failed: {failed})")378 else:379 logger.warning(f"⚠️ [GATEKEEPER] Only {loaded} experts found. Below operational threshold ({MIN_EXPERTS_OPERATIONAL}). Awaiting Forge.")380 381 def reload_brains(self) -> None:382 """383 🚨 HOT-RELOAD: Safely dumps old models from memory and loads the newly forged ones384 without interrupting the live execution thread. Also re-fetches optimal_params.385 """386 logger.info("🔄 [GATEKEEPER] Hot-reloading new Neural Council into RAM...")387 self.params = load_optimal_params()388 self.seq_length = self.params["seq_length"]389 self._load_council()390 391 # --------------------------------------------------------------------------392 # Core Inference Methods393 # --------------------------------------------------------------------------394 def _execute_ensemble_inference(395 self,396 tensor_data: torch.Tensor,397 symbol: str,398 hmm_bull_prior: float = 0.5399 ) -> Tuple[float, float, float]:400 """401 Core inference logic. Returns (consensus_prob, uncertainty, ensemble_value).402 """403 if not self.is_trained or tensor_data is None:404 return hmm_bull_prior, 0.0, 0.0405 406 try:407 regime_idx = infer_regime_idx_from_sequence(tensor_data)408 409 votes = []410 values = []411 412 with torch.no_grad():413 for model in self.council_brains:414 prob, val = model(tensor_data, regime_idx=regime_idx, return_value=True)415 votes.append(prob.item())416 values.append(val.item())417 418 # Consensus Probability (Mean of policy heads)419 consensus_prob = float(np.mean(votes))420 421 # Epistemic Uncertainty (Standard Deviation of policy heads)422 uncertainty = float(np.std(votes))423 424 # Ensemble Value Score (Mean of value heads, used for PSO steering)425 ensemble_value = float(np.mean(values))426 427 # Logging extreme high-conviction setups428 if consensus_prob > 0.80 and uncertainty < 0.05:429 logger.info(f"🎯 [COUNCIL ALIGNMENT] {len(self.council_brains)}/{len(self.council_brains)} Experts agree on {symbol} Breakout. (Conf: {consensus_prob:.2f} | Val: {ensemble_value:.2f})")430 elif consensus_prob < 0.20 and uncertainty < 0.05:431 logger.info(f"🩸 [COUNCIL ALIGNMENT] {len(self.council_brains)}/{len(self.council_brains)} Experts agree on {symbol} Collapse. (Conf: {consensus_prob:.2f} | Val: {ensemble_value:.2f})")432 433 return consensus_prob, uncertainty, ensemble_value434 435 except Exception as e:436 logger.error(f"Gatekeeper Inference Error: {e}")437 return hmm_bull_prior, 0.0, 0.0438 439 # --------------------------------------------------------------------------440 # Public APIs441 # --------------------------------------------------------------------------442 def get_ensemble_probability_from_dict(443 self, 444 symbol: str, 445 tensor_dict: dict, 446 hmm_bull_prior: float = 0.5447 ) -> Tuple[float, float, float]:448 """449 Inference using a single 16D snapshot (alpha_physics tensor_dict).450 The snapshot is tiled backwards to fill the GRU's sequence length requirement.451 """452 tensor_data = build_tensor_from_physics(tensor_dict, seq_length=self.seq_length)453 return self._execute_ensemble_inference(tensor_data, symbol, hmm_bull_prior)454 455 def get_ensemble_probability_from_history(456 self, 457 symbol: str, 458 feature_history: List[List[float]], 459 hmm_bull_prior: float = 0.5460 ) -> Tuple[float, float, float]:461 """462 Inference using a true temporal rolling history of 16D snapshots.463 This provides the highest fidelity to the GRU as it captures the actual time-series evolution.464 """465 tensor_data = build_tensor_from_history(feature_history, seq_length=self.seq_length)466 return self._execute_ensemble_inference(tensor_data, symbol, hmm_bull_prior)467