TwoFootBot-rust/src/policy.rs

use std::path::Path;
use anyhow::{Context, Result, ensure};
use tract::prelude::*;

/// 运行策略模型：输入 36 维观测值，返回 8 维动作输出。
///
/// - 每次调用都会重新加载并准备模型（简单、无状态，方便一行调用）。
/// - 模型路径固定为 `CARGO_MANIFEST_DIR/model/policy.onnx`。
/// - 使用 Metal 运行时。
pub fn policy(observation: [f32; 36]) -> Result<[f32; 8]> {
    let model_path = Path::new(env!("CARGO_MANIFEST_DIR")).join("model/policy.onnx");

    // 加载模型
    let model = tract::onnx()?
        .load(&model_path)
        .with_context(|| format!("failed to load ONNX model at {}", model_path.display()))?
        .into_model()?;

    // 校验输入数量与类型
    ensure!(model.input_count()? == 1, "policy model must have exactly one input");
    let input_fact = model.input_fact(0)?;
    ensure!(
        input_fact.datum_type()? == DatumType::F32,
        "policy model input must be f32, got {:?}",
        input_fact.datum_type()?
    );

    // 读取固定输入形状
    let input_shape = (0..input_fact.rank()?)
        .map(|axis| {
            let dim = input_fact
                .dim(axis)?
                .to_int64()
                .with_context(|| format!("input dimension {axis} must be fixed"))?;
            usize::try_from(dim)
                .with_context(|| format!("input dimension {axis} must be non-negative"))
        })
        .collect::<Result<Vec<_>>>()?;

    // 校验 shape 与 observation 长度匹配
    let input_len: usize = input_shape.iter().product();
    ensure!(
        input_len == observation.len(),
        "expected {input_len} observation values for input shape {input_shape:?}, got {}",
        observation.len()
    );

    // 构造输入张量并推理
    let input_tensor = Tensor::from_slice(&input_shape, &observation)?;
    let runtime = tract::runtime_for_name("metal")?;
    let runnable = runtime.prepare(model)?;
    let outputs = runnable.run([input_tensor])?;

    // 取第一个输出并转成 [f32; 8]
    let output = outputs
        .first()
        .context("policy model produced no output")?;
    let values = output.as_slice::<f32>()?;
    ensure!(
        values.len() == 8,
        "expected 8 output values, got {}",
        values.len()
    );

    let mut result = [0.0f32; 8];
    result.copy_from_slice(values);
    result
        .iter_mut()
        .for_each(|value| *value = (*value * 0.5).clamp(-1.57, 1.57));
    Ok(result)
}