koios-model-utils 1.0.0

Utilities for exporting and annotating ONNX models for the Koios IoT platform

koios-model-utils

Utilities for exporting and annotating ONNX models for the Koios IoT platform.

This package helps data scientists prepare trained ML models (especially Stable-Baselines3 reinforcement learning models) for deployment on Koios. It handles ONNX export, metadata embedding, and normalization parameter annotation so the Koios Predict Engine can run inference correctly.

Installation

# Metadata embedding only (no torch/SB3 dependency)
pip install koios-model-utils

# Full SB3 export pipeline
pip install koios-model-utils[sb3]

Quick Start

Embed metadata into an existing ONNX model

import onnx
from koios_model_utils import (
    InputBinding, OutputBinding, TrainingMeta,
    embed_koios_metadata,
)

model = onnx.load("my_model.onnx")

inputs = [
    InputBinding(name="temperature", description="Tank temperature (C)"),
    InputBinding(name="pressure", description="Vessel pressure (kPa)"),
]
outputs = [
    OutputBinding(
        name="valve_position",
        range_min=0.0, range_max=100.0,
        normalization_type="symmetric",
        normalization_source="custom",
        custom_minimum=0.0, custom_maximum=100.0,
        clamp_output=True,
    ),
]
training = TrainingMeta(
    scenario_name="tank_temperature",
    algorithm="PPO",
    obs_depth=5,
    sample_rate=1.0,
)

embed_koios_metadata(model, inputs=inputs, outputs=outputs, training=training)
onnx.save(model, "my_model_koios.onnx")

Export an SB3 model to ONNX

from koios_model_utils import (
    Algorithm,
    TrainingMeta,
    build_input_bindings,
    build_output_bindings,
    export_onnx,
)

inputs = build_input_bindings(
    ["temperature", "level", "flow_rate"],
    descriptions={"temperature": "Tank temp (C)", "level": "Tank level (%)"},
)
outputs = build_output_bindings(
    ["valve_position"],
    action_ranges={"valve_position": (0.0, 100.0)},
)

export_onnx(
    "runs/best_model.zip",
    inputs=inputs,
    outputs=outputs,
    training=TrainingMeta(
        scenario_name="tank_temperature",
        algorithm=Algorithm.PPO,
        obs_depth=5,
    ),
)

CLI

# Export with auto-detected settings
koios-export runs/best_model.zip

# Specify output path and opset version
koios-export runs/best_model.zip --output model.onnx --opset 18

SB3 Training Callbacks

from koios_model_utils import ObsRangeCallback, VecNormalizeSyncCallback

# Track observation ranges during training
obs_range_cb = ObsRangeCallback(run_dir, verbose=1)

# Sync VecNormalize stats with best model checkpoints
vec_sync_cb = VecNormalizeSyncCallback(eval_callback, verbose=1)

API Reference

Metadata Types

Class	Description
InputBinding	Describes one input (observation) binding — name, normalization, failure bounds
OutputBinding	Describes one output (action) binding — name, range, normalization, clamping
TrainingMeta	Model-level metadata — algorithm, obs_depth, sample_rate, model_type

Functions

Function	Description
embed_koios_metadata()	Embed koios.training and koios.bindings metadata into an ONNX model
build_input_bindings()	Build input bindings from node names + optional normalization stats
build_output_bindings()	Build output bindings from action names + ranges
export_onnx()	Full SB3-to-ONNX export pipeline (requires [sb3] extra)
detect_algorithm()	Detect SB3 algorithm (PPO, SAC, TD3, etc.) from a saved .zip model

Callbacks (requires [sb3] extra)

Callback	Description
ObsRangeCallback	Track per-feature min/max of observations during training
VecNormalizeSyncCallback	Save VecNormalize stats alongside best model checkpoints

Development

# Install with all dependencies
pip install -e ".[sb3,dev]"

# Run tests
make test

# Lint
make lint

License

MIT License - see LICENSE for details.