policy is ok

.gitignore

@@ -0,0 +1 @@
+/target

.idea/.gitignore

@@ -0,0 +1,10 @@
+# 默认忽略的文件
+/shelf/
+/workspace.xml
+# 基于编辑器的 HTTP 客户端请求
+/httpRequests/
+# 已忽略包含查询文件的默认文件夹
+/queries/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/TowFootBot-rust.iml

@@ -0,0 +1,11 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="EMPTY_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <sourceFolder url="file://$MODULE_DIR$/src" isTestSource="false" />
+      <excludeFolder url="file://$MODULE_DIR$/target" />
+    </content>
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/TowFootBot-rust.iml" filepath="$PROJECT_DIR$/.idea/TowFootBot-rust.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "TowFootBot-rust"
+version = "0.1.0"
+edition = "2024"
+[dependencies]
+tract = "0.23"
+tract-metal = "0.23"
+ndarray = "0.16"
+anyhow = "1.0"

src/main.rs

@@ -0,0 +1,31 @@
+mod policy;
+use anyhow::{Context, Result};
+// fn main() -> Result<()> {
+//     //循环
+//     for _ in 0..10 {
+//         let action = policy::policy([0.0; 36])?;
+//         println!("模型输出: {action:?}");
+//     }
+//     Ok(())
+// }
+
+use std::time::Instant;
+
+fn main() -> Result<()> {
+    const N: u32 = 1000;
+    let mut total_infer = std::time::Duration::ZERO;
+
+    for _ in 0..N {
+        let start = Instant::now();
+        let action = policy::policy([0.0; 36])?;
+        total_infer += start.elapsed();
+        println!("模型输出: {action:?}");
+    }
+
+    println!(
+        "平均耗时: {:?}",
+        total_infer / N
+    );
+    println!("采样次数: {N}");
+    Ok(())
+}

src/policy.rs

@@ -0,0 +1,68 @@
+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)
+}