botstadia-interface 0.2.0

//! The policy shell's ONNX contract, version 1: how a trained policy's model
//! must look so one generic program runs it in every arena.
//!
//! - **Inputs:** one per declared observation, named exactly as declared,
//!   shape `[1, *shape]`, `f32`. A `u8` value is fed as `f32 / 255`. A policy
//!   may use a subset of the observations.
//! - **Memory:** an optional `state_in: f32[1, S]` → `state_out: f32[1, S]`
//!   pair, `S` at most [`MAX_STATE`]. Zeros at the start; each turn's
//!   `state_out` is the next turn's `state_in`.
//! - **Continuous actions** (declared `f32`): one output per declared action,
//!   named as declared, shape `[1, *shape]`, in `[-1, 1]`. The shell maps each
//!   element with [`to_bounds`]: clamp to `[-1, 1]` (NaN is 0, the midpoint),
//!   then affinely onto the element's declared bounds.
//! - **Discrete actions** (declared `u8`): one output per declared action,
//!   named as declared, of logits shaped `[1, *shape, n]`, where every element
//!   shares one range `[min, max]` and `n = max - min + 1`. The shell takes the
//!   argmax of each element's logits ([`argmax`]) and adds `min`.
//! - **Metadata:** [`META_CONTRACT`], [`META_ARENA`], [`META_PAYLOAD_SCHEMA`],
//!   [`META_TRAIN_VERSION`] and [`META_OPSET`], checked against the manifest.
//! - An input named [`NOISE`] is reserved.
//!
//! `botstadia-train` implements the same map; the shared golden file
//! (`policy-golden.json`, beside this crate's tests) pins both.

/// The golden cases pinning [`to_bounds`], [`feed_u8`] and [`argmax`], which
/// `botstadia-train` checks its own implementation against.
pub const GOLDEN: &str = include_str!("../tests/policy-golden.json");

/// This contract's version, as the `botstadia.contract` metadata states it.
pub const CONTRACT_VERSION: u32 = 1;

pub const META_CONTRACT: &str = "botstadia.contract";
pub const META_ARENA: &str = "botstadia.arena";
pub const META_PAYLOAD_SCHEMA: &str = "botstadia.payload_schema";
pub const META_TRAIN_VERSION: &str = "botstadia.train_version";
pub const META_OPSET: &str = "botstadia.opset";

/// The memory pair's names.
pub const STATE_IN: &str = "state_in";
pub const STATE_OUT: &str = "state_out";
/// Reserved for sampled actions later.
pub const NOISE: &str = "noise";
/// The largest memory, in floats.
pub const MAX_STATE: u32 = 1024;
/// The opset the export helper writes.
pub const EXPORT_OPSET: i64 = 18;

/// One element of a continuous action: a policy output in `[-1, 1]` mapped
/// onto the declared `[min, max]`. Out-of-range outputs are clamped and NaN
/// is the midpoint, so the result is always valid.
pub fn to_bounds(value: f32, min: f32, max: f32) -> f32 {
    let unit = if value.is_nan() {
        0.0
    } else {
        value.clamp(-1.0, 1.0)
    };
    let mapped = min + (unit + 1.0) * 0.5 * (max - min);
    // Rounding can step a hair outside the bounds; the host rejects that.
    mapped.clamp(min, max)
}

/// The inverse of [`to_bounds`], for training: a declared value as the
/// policy's `[-1, 1]`.
pub fn from_bounds(value: f32, min: f32, max: f32) -> f32 {
    if max <= min {
        return 0.0;
    }
    ((value - min) / (max - min) * 2.0 - 1.0).clamp(-1.0, 1.0)
}

/// An observation element as the policy sees it.
pub fn feed_u8(value: u8) -> f32 {
    f32::from(value) / 255.0
}

/// The index of the largest logit; NaN never wins, and ties go to the first.
pub fn argmax(logits: &[f32]) -> usize {
    let mut best = 0;
    let mut best_value = f32::NEG_INFINITY;
    for (index, value) in logits.iter().enumerate() {
        if *value > best_value {
            best = index;
            best_value = *value;
        }
    }
    best
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The golden cases `botstadia-train` checks too.
    #[derive(serde::Deserialize)]
    struct Golden {
        #[serde(rename = "_doc")]
        _doc: String,
        to_bounds: Vec<(f32, f32, f32, f32)>,
        feed_u8: Vec<(u8, f32)>,
        argmax: Vec<(Vec<f32>, usize)>,
    }

    #[test]
    fn matches_the_shared_golden_file() {
        let golden: Golden = serde_json::from_str(GOLDEN).unwrap();
        for (value, min, max, expected) in golden.to_bounds {
            assert_eq!(
                to_bounds(value, min, max),
                expected,
                "{value} in {min}..{max}"
            );
        }
        for (value, expected) in golden.feed_u8 {
            assert_eq!(feed_u8(value), expected);
        }
        for (logits, expected) in golden.argmax {
            assert_eq!(argmax(&logits), expected);
        }
    }

    #[test]
    fn nan_is_the_midpoint_and_the_map_inverts() {
        assert_eq!(to_bounds(f32::NAN, -2.0, 4.0), 1.0);
        for value in [-1.0, -0.25, 0.0, 0.6, 1.0] {
            let there = to_bounds(value, -2.5, 2.88);
            assert!((from_bounds(there, -2.5, 2.88) - value).abs() < 1e-6);
        }
    }
}