botstadia_interface/
policy.rs

1//! The policy shell's ONNX contract, version 1: how a trained policy's model
2//! must look so one generic program runs it in every arena.
3//!
4//! - **Inputs:** one per declared observation, named exactly as declared,
5//!   shape `[1, *shape]`, `f32`. A `u8` value is fed as `f32 / 255`. A policy
6//!   may use a subset of the observations.
7//! - **Memory:** an optional `state_in: f32[1, S]` → `state_out: f32[1, S]`
8//!   pair, `S` at most [`MAX_STATE`]. Zeros at the start; each turn's
9//!   `state_out` is the next turn's `state_in`.
10//! - **Continuous actions** (declared `f32`): one output per declared action,
11//!   named as declared, shape `[1, *shape]`, in `[-1, 1]`. The shell maps each
12//!   element with [`to_bounds`]: clamp to `[-1, 1]` (NaN is 0, the midpoint),
13//!   then affinely onto the element's declared bounds.
14//! - **Discrete actions** (declared `u8`): one output per declared action,
15//!   named as declared, of logits shaped `[1, *shape, n]`, where every element
16//!   shares one range `[min, max]` and `n = max - min + 1`. The shell takes the
17//!   argmax of each element's logits ([`argmax`]) and adds `min`.
18//! - **Metadata:** [`META_CONTRACT`], [`META_ARENA`], [`META_PAYLOAD_SCHEMA`],
19//!   [`META_TRAIN_VERSION`] and [`META_OPSET`], checked against the manifest.
20//! - An input named [`NOISE`] is reserved.
21//!
22//! `botstadia-train` implements the same map; the shared golden file
23//! (`policy-golden.json`, beside this crate's tests) pins both.
24
25/// The golden cases pinning [`to_bounds`], [`feed_u8`] and [`argmax`], which
26/// `botstadia-train` checks its own implementation against.
27pub const GOLDEN: &str = include_str!("../tests/policy-golden.json");
28
29/// This contract's version, as the `botstadia.contract` metadata states it.
30pub const CONTRACT_VERSION: u32 = 1;
31
32pub const META_CONTRACT: &str = "botstadia.contract";
33pub const META_ARENA: &str = "botstadia.arena";
34pub const META_PAYLOAD_SCHEMA: &str = "botstadia.payload_schema";
35pub const META_TRAIN_VERSION: &str = "botstadia.train_version";
36pub const META_OPSET: &str = "botstadia.opset";
37
38/// The memory pair's names.
39pub const STATE_IN: &str = "state_in";
40pub const STATE_OUT: &str = "state_out";
41/// Reserved for sampled actions later.
42pub const NOISE: &str = "noise";
43/// The largest memory, in floats.
44pub const MAX_STATE: u32 = 1024;
45/// The opset the export helper writes.
46pub const EXPORT_OPSET: i64 = 18;
47
48/// One element of a continuous action: a policy output in `[-1, 1]` mapped
49/// onto the declared `[min, max]`. Out-of-range outputs are clamped and NaN
50/// is the midpoint, so the result is always valid.
51pub fn to_bounds(value: f32, min: f32, max: f32) -> f32 {
52    let unit = if value.is_nan() {
53        0.0
54    } else {
55        value.clamp(-1.0, 1.0)
56    };
57    let mapped = min + (unit + 1.0) * 0.5 * (max - min);
58    // Rounding can step a hair outside the bounds; the host rejects that.
59    mapped.clamp(min, max)
60}
61
62/// The inverse of [`to_bounds`], for training: a declared value as the
63/// policy's `[-1, 1]`.
64pub fn from_bounds(value: f32, min: f32, max: f32) -> f32 {
65    if max <= min {
66        return 0.0;
67    }
68    ((value - min) / (max - min) * 2.0 - 1.0).clamp(-1.0, 1.0)
69}
70
71/// An observation element as the policy sees it.
72pub fn feed_u8(value: u8) -> f32 {
73    f32::from(value) / 255.0
74}
75
76/// The index of the largest logit; NaN never wins, and ties go to the first.
77pub fn argmax(logits: &[f32]) -> usize {
78    let mut best = 0;
79    let mut best_value = f32::NEG_INFINITY;
80    for (index, value) in logits.iter().enumerate() {
81        if *value > best_value {
82            best = index;
83            best_value = *value;
84        }
85    }
86    best
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    /// The golden cases `botstadia-train` checks too.
94    #[derive(serde::Deserialize)]
95    struct Golden {
96        #[serde(rename = "_doc")]
97        _doc: String,
98        to_bounds: Vec<(f32, f32, f32, f32)>,
99        feed_u8: Vec<(u8, f32)>,
100        argmax: Vec<(Vec<f32>, usize)>,
101    }
102
103    #[test]
104    fn matches_the_shared_golden_file() {
105        let golden: Golden = serde_json::from_str(GOLDEN).unwrap();
106        for (value, min, max, expected) in golden.to_bounds {
107            assert_eq!(
108                to_bounds(value, min, max),
109                expected,
110                "{value} in {min}..{max}"
111            );
112        }
113        for (value, expected) in golden.feed_u8 {
114            assert_eq!(feed_u8(value), expected);
115        }
116        for (logits, expected) in golden.argmax {
117            assert_eq!(argmax(&logits), expected);
118        }
119    }
120
121    #[test]
122    fn nan_is_the_midpoint_and_the_map_inverts() {
123        assert_eq!(to_bounds(f32::NAN, -2.0, 4.0), 1.0);
124        for value in [-1.0, -0.25, 0.0, 0.6, 1.0] {
125            let there = to_bounds(value, -2.5, 2.88);
126            assert!((from_bounds(there, -2.5, 2.88) - value).abs() < 1e-6);
127        }
128    }
129}
botstadia_interface/policy.rs

botstadia_interface/
policy.rs
