68 lines
2.3 KiB
Rust
68 lines
2.3 KiB
Rust
|
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);
|
||
|
|
Ok(result)
|
||
|
|
}
|