""" mnn_export.py ============= """ from __future__ import annotations import argparse import hashlib import logging import os import sys from pathlib import Path from typing import Dict, List, Optional, Tuple import zone_observation as _zo assert _zo.SCHEMA_VERSION == 3, ( f"mnn_export: zone_observation schema mismatch " f"(expected 3, got {_zo.SCHEMA_VERSION})" ) from zone_observation import ForecastConfig try: from weather_forecast_env import make_weather_env from gru_weather_policy import GRUWeatherFeaturesExtractor, create_gru_weather_policy_kwargs _ML_AVAILABLE = True except ImportError: _ML_AVAILABLE = False GRUWeatherFeaturesExtractor = None # type: ignore[assignment,misc] create_gru_weather_policy_kwargs = None # type: ignore[assignment] make_weather_env = None # type: ignore[assignment] logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Hardware capability reporter # --------------------------------------------------------------------------- def report_edge_capability() -> str: """ Detect GPU/compute backend on the current device. Returns one of: 'VULKAN', 'OPENCL', 'CPU_ONLY'. """ import subprocess try: out = subprocess.run( ["vulkaninfo", "--summary"], capture_output=True, text=True, timeout=5, ).stdout if "Vulkan" in out: logger.info("Edge capability: Vulkan detected -> MNN Vulkan backend") return "VULKAN" except (FileNotFoundError, subprocess.TimeoutExpired): pass try: out = subprocess.run( ["clinfo"], capture_output=True, text=True, timeout=5, ).stdout if "Mali" in out: if "OpenCL 2" in out or "OpenCL 3" in out: logger.info("Edge capability: Mali + OpenCL 2/3 -> MNN OpenCL backend") return "OPENCL" if "Mali-450" in out or "Utgard" in out: logger.warning( "Edge capability: Mali-450 (Utgard) detected. " "No OpenCL / Vulkan support. Forcing CPU fallback. " "The quantized .mnn file will still run — just slower." ) return "CPU_ONLY" except (FileNotFoundError, subprocess.TimeoutExpired): pass logger.info("Edge capability: GPU info unavailable -> defaulting to CPU_ONLY") return "CPU_ONLY" # --------------------------------------------------------------------------- # StatelessInferenceWrapper — explicit hidden state I/O for GRU tracing # --------------------------------------------------------------------------- class StatelessInferenceWrapper: """ Wraps the PPO actor so hidden state is explicit I/O rather than Python state. WHY THIS EXISTS: torch.jit.trace records a single execution path. GRUWeatherFeaturesExtractor keeps _hidden as a Python object (None on first call, Tensor thereafter). This branching logic is invisible to the tracer — the resulting graph would always reinitialise hidden state, silently breaking temporal belief propagation on edge. This wrapper eliminates the branch: hidden_state is accepted as an explicit input tensor and returned as an explicit output. The edge runtime manages hidden state externally between steps. INPUTS (fixed float32): obs_tensors: list of observation tensors in canonical key order hidden_in: [1, 1, hidden_size] float32 GRU hidden state OUTPUTS: action_logits: [1, n_actions] float32 — apply mask + argmax on edge hidden_out: [1, 1, hidden_size] float32 — feed back next step The interface is identical to the previous DQN export from the edge's perspective: a vector of per-action scores, a mask, and persistent hidden state. No changes needed to edge runtime code. """ def __init__( self, features_extractor, mlp_extractor, action_net, hidden_size: int, obs_keys: List[str], ): self.features_extractor = features_extractor self.mlp_extractor = mlp_extractor self.action_net = action_net self.hidden_size = hidden_size self.obs_keys = obs_keys def forward( self, obs_tensors: List, # one tensor per obs_key, in canonical order hidden_in, # [1, 1, hidden_size] ) -> Tuple: """Pure function — no Python-object state. Safe to trace.""" # Rebuild obs dict from positional tensors (tracing-safe) obs = {k: obs_tensors[i] for i, k in enumerate(self.obs_keys)} # Inject external hidden state into the features extractor self.features_extractor.set_hidden(hidden_in) # Extract features (MLP/conv + GRU step) features = self.features_extractor(obs) # Actor path only — discard value/critic at export time latent_pi = self.mlp_extractor.forward_actor(features) # Action logits [1, n_actions] action_logits = self.action_net(latent_pi) # Return updated hidden state for the edge runtime to store hidden_out = self.features_extractor.get_hidden() return action_logits, hidden_out # --------------------------------------------------------------------------- # Validation helpers # --------------------------------------------------------------------------- def _validate_output_path(path_str: str) -> Path: p = Path(path_str).resolve() if p.suffix != ".mnn": raise ValueError( f"Output path must end with .mnn, got: {path_str!r}" ) p.parent.mkdir(parents=True, exist_ok=True) return p def _sha256_file(path: Path, chunk_size: int = 1 << 20) -> str: h = hashlib.sha256() with open(path, "rb") as f: while chunk := f.read(chunk_size): h.update(chunk) return h.hexdigest() # --------------------------------------------------------------------------- # Calibration data collection # --------------------------------------------------------------------------- def _collect_calibration_obs( env, n_episodes: int = 100, obs_keys: Optional[List[str]] = None, ) -> List[Dict]: """ Run random episodes to collect representative observations for PTQ calibration. 100–200 episodes is sufficient for most RL policies. """ import numpy as np logger.info("Collecting calibration data (%d episodes)...", n_episodes) samples = [] for ep in range(n_episodes): obs, _ = env.reset() done = False steps = 0 while not done and steps < 50: sample = {k: v for k, v in obs.items() if k != "action_mask"} samples.append(sample) action = env.action_space.sample() obs, _, terminated, truncated, _ = env.step(action) done = terminated or truncated steps += 1 if (ep + 1) % 20 == 0: logger.info(" Calibration: %d/%d episodes", ep + 1, n_episodes) logger.info("Collected %d calibration samples", len(samples)) return samples # --------------------------------------------------------------------------- # ONNX export # --------------------------------------------------------------------------- def _export_onnx( wrapper: StatelessInferenceWrapper, dummy_obs: Dict, hidden_size: int, onnx_path: Path, export_keys: List[str], ) -> None: """ Export the stateless actor wrapper to ONNX. action_mask is excluded from the ONNX graph — it is applied by the edge runtime after receiving action logits (mask → argmax protocol). hidden_state is explicit I/O for temporal belief propagation. """ import torch obs_tensors = [dummy_obs[k].float() for k in export_keys] dummy_hidden = torch.zeros(1, 1, hidden_size) input_names = export_keys + ["hidden_in"] output_names = ["action_logits", "hidden_out"] dynamic_axes: Dict[str, Dict[int, str]] = {k: {0: "batch"} for k in export_keys} dynamic_axes["hidden_in"] = {1: "batch"} dynamic_axes["action_logits"] = {0: "batch"} dynamic_axes["hidden_out"] = {1: "batch"} logger.info("Exporting to ONNX (opset 17): %s", onnx_path) logger.info(" Observation inputs: %s", export_keys) logger.info(" Outputs: %s", output_names) with torch.no_grad(): torch.onnx.export( wrapper, args=(obs_tensors, dummy_hidden), f=str(onnx_path), opset_version=17, input_names=input_names, output_names=output_names, dynamic_axes=dynamic_axes, ) logger.info( "ONNX saved: %s (%.1f MB)", onnx_path, onnx_path.stat().st_size / 1e6, ) # --------------------------------------------------------------------------- # MNN conversion # --------------------------------------------------------------------------- def _convert_to_mnn( onnx_path: Path, mnn_path: Path, quantize: str, calibration_samples: Optional[List[Dict]] = None, ) -> None: """ Convert ONNX to MNN using the MNN Python API, with CLI fallback. quantize options: 'int8' — weight-only INT8 (no calibration needed, recommended) 'fp16' — FP16 half-precision (higher accuracy, ~2x size reduction) 'none' — FP32 (largest, highest accuracy, use for debugging) """ logger.info( "Converting to MNN quantize=%s target=%s", quantize, mnn_path ) converted = False try: from MNN.tools import mnnconvert as _mnnconvert args = { "modelFile": str(onnx_path), "MNNModel": str(mnn_path), "framework": "ONNX", "bizCode": "weather_rl_v1", } if quantize == "int8": args["weightQuantBits"] = 8 elif quantize == "fp16": args["fp16"] = True _mnnconvert.convert(args) converted = True logger.info("MNN conversion via Python API: OK") except Exception as api_err: logger.warning("MNN Python API failed (%s) — trying CLI fallback", api_err) if not converted: import subprocess, shutil cli = shutil.which("mnnconvert") if cli is None: raise RuntimeError( "mnnconvert not found on PATH and MNN Python API failed.\n" "Build from: https://github.com/GeniusVentures/MNN\n" "Or install: pip install MNN" ) cmd = [cli, "-f", "ONNX", "--modelFile", str(onnx_path), "--MNNModel", str(mnn_path)] if quantize == "int8": cmd += ["--weightQuantBits", "8"] elif quantize == "fp16": cmd += ["--fp16"] logger.info("MNN CLI: %s", " ".join(cmd)) result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise RuntimeError( f"mnnconvert CLI failed (rc={result.returncode}):\n" f"stdout: {result.stdout}\nstderr: {result.stderr}" ) if not mnn_path.exists(): raise RuntimeError( f"MNN conversion reported success but {mnn_path} was not created." ) # --------------------------------------------------------------------------- # Main export function # --------------------------------------------------------------------------- def export_to_mnn( checkpoint_path: str, output_mnn: str = "weather_rl_model.mnn", quantize: str = "int8", calibration_episodes: int = 0, hidden_size: int = 64, n_zones: int = 4, keep_onnx: bool = False, ) -> Path: """ Export a trained MaskablePPO checkpoint to a quantized .mnn for edge deployment. Args: checkpoint_path: Path to the trained .zip checkpoint. output_mnn: Output .mnn file path (must end with .mnn). quantize: 'int8' (default), 'fp16', or 'none'. calibration_episodes: Episodes for PTQ calibration (0 = weight-only). hidden_size: GRU hidden size used during training (default 64). n_zones: Number of zones the checkpoint was trained with. Must match the curriculum phase: normal=2, monsoon/drought=3, heatwave/humidity=4 (default 4). ForecastConfig() defaults to n_zones=1, which is wrong for any multi-zone checkpoint — always pass the value that matches the training phase explicitly. keep_onnx: If True, keep the intermediate .onnx file. Returns: Path to the created .mnn file. """ import torch from sb3_contrib import MaskablePPO ckpt = Path(checkpoint_path) if not ckpt.exists(): raise FileNotFoundError(f"Checkpoint not found: {ckpt}") if quantize not in ("int8", "fp16", "none"): raise ValueError(f"quantize must be 'int8', 'fp16', or 'none', got {quantize!r}") mnn_path = _validate_output_path(output_mnn) onnx_path = mnn_path.with_suffix(".onnx") backend = report_edge_capability() # --- Load MaskablePPO checkpoint --- logger.info("Loading checkpoint: %s", ckpt) # FIX: use the n_zones that matches the training phase, not ForecastConfig() # default of n_zones=1. The dummy environment produced by env.reset() is # used only to build dummy_obs for ONNX tracing; its tensor shapes must # match those of the loaded policy or the traced graph will have wrong # input shapes and be incompatible with edge_wrapper.cpp (N_ZONES=4). if n_zones < 1: raise ValueError(f"n_zones must be >= 1, got {n_zones}") config = ForecastConfig(n_zones=n_zones, horizon_days=30) env = make_weather_env(config) logger.info("Export env: n_zones=%d horizon_days=30", n_zones) custom_objects = {} if GRUWeatherFeaturesExtractor is not None: # FIX (schema v3 / train_kaggle compat): pass ONLY the class, never a # hardcoded features_extractor_kwargs. SB3 restores the extractor's # kwargs from the checkpoint's saved policy_kwargs; anything passed # here OVERRIDES them. The previous version forced # spatial_output_size=8 and (by omission) basin_context_hidden=8, # which mismatches train_kaggle.py checkpoints (trained with # spatial_output_size=12, basin_context_hidden=12) and made # MaskablePPO.load fail with a state_dict size mismatch -- i.e. the # export path could not load the project's own training output. custom_objects = { "features_extractor_class": GRUWeatherFeaturesExtractor, } model = MaskablePPO.load( str(ckpt), env=env, device="cpu", custom_objects=custom_objects if custom_objects else None, ) model.policy.eval() # Verify the loaded policy has the expected actor components. # This catches mismatches between the checkpoint and the export path # (e.g. a checkpoint saved with a custom policy that removed mlp_extractor). policy = model.policy assert hasattr(policy, "mlp_extractor") and hasattr(policy, "action_net"), ( f"Loaded policy is missing expected actor components. " f"Got attributes: {[a for a in dir(policy) if not a.startswith('_')]}" ) assert hasattr(policy, "features_extractor"), ( "Loaded policy is missing features_extractor." ) logger.info("Checkpoint loaded: %s", ckpt.name) # Trust the checkpoint over the CLI flag for the GRU hidden size: the # wrapper's hidden_in/hidden_out tensor shape must equal the trained # GRU's, and a stale --hidden-size default would silently trace a # wrong-shaped graph. _ckpt_hidden = getattr(policy.features_extractor, "hidden_size", None) if _ckpt_hidden is not None and _ckpt_hidden != hidden_size: logger.warning( "Overriding --hidden-size=%d with checkpoint's hidden_size=%d", hidden_size, _ckpt_hidden, ) hidden_size = _ckpt_hidden # --- Build stateless actor wrapper --- # We export the actor path only: # features_extractor → mlp_extractor.forward_actor → action_net # The critic (value_net) is discarded — not needed at inference time. features_extractor = policy.features_extractor mlp_extractor = policy.mlp_extractor action_net = policy.action_net obs_sample, _ = env.reset() obs_keys_all = sorted(obs_sample.keys()) export_keys = [k for k in obs_keys_all if k != "action_mask"] wrapper = StatelessInferenceWrapper( features_extractor=features_extractor, mlp_extractor=mlp_extractor, action_net=action_net, hidden_size=hidden_size, obs_keys=export_keys, ) # --- Build dummy input --- dummy_obs: Dict[str, "torch.Tensor"] = {} for k, v in obs_sample.items(): t = torch.from_numpy(v).unsqueeze(0) dummy_obs[k] = t.float() if k != "action_mask" else t # --- Optional calibration --- calibration_samples = None if calibration_episodes > 0 and quantize == "int8": calibration_samples = _collect_calibration_obs( env, n_episodes=calibration_episodes ) # --- ONNX export --- _export_onnx(wrapper, dummy_obs, hidden_size, onnx_path, export_keys) # --- MNN conversion --- try: _convert_to_mnn(onnx_path, mnn_path, quantize, calibration_samples) finally: if onnx_path.exists() and not keep_onnx: onnx_path.unlink() logger.info("Removed intermediate ONNX: %s", onnx_path.name) if not mnn_path.exists(): raise RuntimeError( f"Export appeared to succeed but {mnn_path} was not created." ) size_mb = mnn_path.stat().st_size / (1024 * 1024) sha = _sha256_file(mnn_path) logger.info( "MNN export complete: %s (%.1f MB) SHA256: %s", mnn_path, size_mb, sha ) logger.info("Edge backend detected: %s", backend) print(f""" ╔══════════════════════════════════════════════════════════╗ ║ Weather RL Model — Edge Deployment Manifest ║ ╠══════════════════════════════════════════════════════════╣ ║ Model: {mnn_path.name:<44} ║ ║ Size: {f'{size_mb:.1f} MB':<44} ║ ║ Quantize: {quantize:<44} ║ ║ Backend: {backend:<44} ║ ║ SHA256: {sha[:44]} ║ ║ {sha[44:]} ║ ╠══════════════════════════════════════════════════════════╣ ║ POST-PROCESSING (apply in edge runtime): ║ ║ logits = model.run(obs_without_mask, hidden_in) ║ ║ logits[action_mask == 0] = -1e9 ║ ║ action = argmax(logits) ║ ║ Store hidden_out; pass as hidden_in next step ║ ╚══════════════════════════════════════════════════════════╝ """) return mnn_path # --------------------------------------------------------------------------- # Edge inference protocol (copy into edge runtime documentation) # --------------------------------------------------------------------------- EDGE_INFERENCE_NOTE = """ Edge Runtime Inference Protocol ================================ The exported .mnn model is a stateless actor network. The edge runtime must manage two pieces of state externally: 1. GRU hidden state (temporal belief): - Initialise: hidden = zeros([1, 1, {hidden_size}]) - Each step: action_logits, hidden = model.run(obs_inputs, hidden) - Reset: hidden = zeros([1, 1, {hidden_size}]) at episode start 2. Action mask (zone validity): - The model outputs raw action logits [1, n_zones + 1] - Apply mask BEFORE argmax: action_logits[action_mask == 0] = -1e9 action = argmax(action_logits) - The terminate action (index n_zones) is ALWAYS valid; never mask it. Input tensor order (must match ONNX input_names exactly): {obs_keys_without_mask} (float32) hidden_in (float32, shape [1, 1, H]) Output tensors: action_logits float32 [1, n_zones + 1] raw scores; apply mask + argmax hidden_out float32 [1, 1, H] store and feed back next step Note (schema v3): the observation inputs now include basin_context (float32, shape [1, 4] -- [enso_oni, iod_dmi, itcz_latitude, mslp_anomaly] in that field order, matching zone_observation.BasinContext.to_array()). It sorts FIRST in the alphabetical input order above. edge_wrapper.cpp has been updated to match; any other edge runtime built against the pre-v3 (4-input) interface MUST add this input -- the MNN session will fail or produce garbage logits if basin_context is left unbound. When no basin data is available at the edge, feed the neutral default [0.0, 0.0, 0.0, 1013.25] (same default weather_forecast_env.py uses). """ # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def _parse_args() -> argparse.Namespace: p = argparse.ArgumentParser( description="Export MaskablePPO weather policy to quantized .mnn", formatter_class=argparse.RawDescriptionHelpFormatter, ) p.add_argument("--checkpoint", required=True, help="Path to trained .zip checkpoint") p.add_argument("--output", default="weather_rl_model.mnn", help="Output .mnn path") p.add_argument("--quantize", default="int8", choices=["int8", "fp16", "none"]) p.add_argument("--calibration-episodes", type=int, default=0, help="Episodes for PTQ calibration (0 = weight-only)") p.add_argument("--hidden-size", type=int, default=64) p.add_argument( "--n-zones", type=int, default=4, help=( "Zones the checkpoint was trained with. " "normal=2, monsoon/drought=3, heatwave/humidity=4 (default 4). " "Must match the curriculum phase or the ONNX trace will have wrong " "input shapes and be incompatible with edge_wrapper.cpp." ), ) p.add_argument("--keep-onnx", action="store_true") p.add_argument("--capability", action="store_true", help="Report edge hardware capability and exit") return p.parse_args() def main() -> None: logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", ) args = _parse_args() if args.capability: print(f"Edge backend: {report_edge_capability()}") sys.exit(0) try: export_to_mnn( checkpoint_path=args.checkpoint, output_mnn=args.output, quantize=args.quantize, calibration_episodes=args.calibration_episodes, hidden_size=args.hidden_size, n_zones=args.n_zones, keep_onnx=args.keep_onnx, ) sys.exit(0) except FileNotFoundError as e: logger.error("Checkpoint not found: %s", e) sys.exit(2) except ValueError as e: logger.error("Invalid argument: %s", e) sys.exit(2) except RuntimeError as e: logger.error("Export failed: %s", e) sys.exit(1) except Exception as e: logger.exception("Unexpected error: %s", e) sys.exit(1) # --------------------------------------------------------------------------- # Self-test (no checkpoint or MNN required) # --------------------------------------------------------------------------- if __name__ == "__main__": if "--checkpoint" in sys.argv: main() else: logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s", ) print("=== mnn_export.py self-test (no checkpoint/MNN required) ===\n") failures = [] def _assert(cond: bool, msg: str) -> None: if not cond: failures.append(msg) print(f" FAIL: {msg}") # 1. _validate_output_path rejects non-.mnn extensions try: _validate_output_path("/tmp/model.pkl") _assert(False, "Should have rejected .pkl extension") except ValueError: pass try: p = _validate_output_path("/tmp/test_export.mnn") _assert(p.suffix == ".mnn", "Resolved path should end in .mnn") except Exception as e: _assert(False, f"Valid .mnn path rejected: {e}") print(" _validate_output_path OK") # 2. StatelessInferenceWrapper can be constructed with mock components try: import torch import torch.nn as nn class _FakeExtractor: def __call__(self, obs): return torch.zeros(1, 256) def set_hidden(self, h): self._h = h def get_hidden(self): return getattr(self, '_h', torch.zeros(1,1,64)) class _FakeMLPExtractor(nn.Module): def forward_actor(self, x): return x[:, :128] wrapper = StatelessInferenceWrapper( features_extractor=_FakeExtractor(), mlp_extractor=_FakeMLPExtractor(), action_net=nn.Linear(128, 5), hidden_size=64, obs_keys=["basin_context", "forecast_precip", "forecast_uncertainty", "prior_belief", "zone_belief"], ) _assert(wrapper.hidden_size == 64, "Wrong hidden_size on wrapper") _assert(len(wrapper.obs_keys) == 5, "Wrong obs_keys count") print(" StatelessInferenceWrapper construction OK") except ImportError as e: print(f" StatelessInferenceWrapper: torch not installed, skipped ({e})") # 3. _sha256_file is deterministic import tempfile with tempfile.NamedTemporaryFile(delete=False, suffix=".bin") as f: f.write(b"weather_rl_test" * 1000) tmp = Path(f.name) sha1 = _sha256_file(tmp) sha2 = _sha256_file(tmp) _assert(sha1 == sha2, "SHA256 not deterministic") _assert(len(sha1) == 64, f"SHA256 wrong length: {len(sha1)}") tmp.unlink() print(f" _sha256_file OK sha={sha1[:16]}...") # 4. report_edge_capability returns a known string cap = report_edge_capability() _assert(cap in ("VULKAN", "OPENCL", "CPU_ONLY"), f"Unknown capability: {cap!r}") print(f" report_edge_capability OK backend={cap}") # 5. SCHEMA_VERSION guard _assert(_zo.SCHEMA_VERSION == 3, f"SCHEMA_VERSION guard not working (got {_zo.SCHEMA_VERSION})") print(" SCHEMA_VERSION guard OK") # 6. EDGE_INFERENCE_NOTE is complete _assert(len(EDGE_INFERENCE_NOTE) > 100, "EDGE_INFERENCE_NOTE too short") _assert("hidden_out" in EDGE_INFERENCE_NOTE, "missing hidden_out") _assert("argmax" in EDGE_INFERENCE_NOTE, "missing argmax") _assert("action_logits" in EDGE_INFERENCE_NOTE, "missing action_logits") print(" EDGE_INFERENCE_NOTE present and complete") # 7. export_keys excludes action_mask (schema v3 key set incl. basin_context) sample_keys = ["action_mask", "basin_context", "forecast_precip", "forecast_uncertainty", "prior_belief", "zone_belief"] export = [k for k in sorted(sample_keys) if k != "action_mask"] _assert("action_mask" not in export, "action_mask leaked into export_keys") _assert(len(export) == 5, f"Expected 5 export keys, got {len(export)}") _assert(export[0] == "basin_context", "basin_context should sort first (alphabetical)") print(" export_keys exclusion of action_mask OK") print() if failures: print(f"FAILED {len(failures)} test(s):") for f in failures: print(f" - {f}") sys.exit(1) else: print("All mnn_export self-tests passed.") print() print("To run the real export:") print(" python mnn_export.py --checkpoint final_normal.zip") print(" python mnn_export.py --capability")