diff --git a/.idea/copilotDiffState.xml b/.idea/copilotDiffState.xml
new file mode 100644
index 0000000..c3642ba
--- /dev/null
+++ b/.idea/copilotDiffState.xml
@@ -0,0 +1,27 @@
+
+
+
+
+
+
\ No newline at end of file
diff --git a/Cargo.lock b/Cargo.lock
index e69db54..c115a1b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -8,6 +8,7 @@ version = "0.1.0"
dependencies = [
"anyhow",
"ndarray 0.16.1",
+ "rdev",
"tract",
"tract-metal",
]
@@ -176,6 +177,12 @@ dependencies = [
"serde",
]
+[[package]]
+name = "bitflags"
+version = "1.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"
+
[[package]]
name = "bitflags"
version = "2.13.2"
@@ -324,6 +331,21 @@ version = "1.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1c133bc6a41be0d194c306b5506d15e6feeea7b1d6604bd3f8310dfb2ca96486"
+[[package]]
+name = "cocoa"
+version = "0.22.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "667fdc068627a2816b9ff831201dd9864249d6ee8d190b9532357f1fc0f61ea7"
+dependencies = [
+ "bitflags 1.3.2",
+ "block",
+ "core-foundation 0.9.4",
+ "core-graphics 0.21.0",
+ "foreign-types 0.3.2",
+ "libc",
+ "objc",
+]
+
[[package]]
name = "colorchoice"
version = "1.0.5"
@@ -336,30 +358,80 @@ version = "1.0.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4e18bf7a165bf7028fde98609a0f1e8f7498d762a212598e6c891f6893556ec"
+[[package]]
+name = "core-foundation"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "57d24c7a13c43e870e37c1556b74555437870a04514f7685f5b354e090567171"
+dependencies = [
+ "core-foundation-sys 0.7.0",
+ "libc",
+]
+
+[[package]]
+name = "core-foundation"
+version = "0.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f"
+dependencies = [
+ "core-foundation-sys 0.8.7",
+ "libc",
+]
+
[[package]]
name = "core-foundation"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6"
dependencies = [
- "core-foundation-sys",
+ "core-foundation-sys 0.8.7",
"libc",
]
+[[package]]
+name = "core-foundation-sys"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3a71ab494c0b5b860bdc8407ae08978052417070c2ced38573a9157ad75b8ac"
+
[[package]]
name = "core-foundation-sys"
version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
+[[package]]
+name = "core-graphics"
+version = "0.19.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b3889374e6ea6ab25dba90bb5d96202f61108058361f6dc72e8b03e6f8bbe923"
+dependencies = [
+ "bitflags 1.3.2",
+ "core-foundation 0.7.0",
+ "foreign-types 0.3.2",
+ "libc",
+]
+
+[[package]]
+name = "core-graphics"
+version = "0.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "52a67c4378cf203eace8fb6567847eb641fd6ff933c1145a115c6ee820ebb978"
+dependencies = [
+ "bitflags 1.3.2",
+ "core-foundation 0.9.4",
+ "foreign-types 0.3.2",
+ "libc",
+]
+
[[package]]
name = "core-graphics-types"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb"
dependencies = [
- "bitflags",
- "core-foundation",
+ "bitflags 2.13.2",
+ "core-foundation 0.10.1",
"libc",
]
@@ -534,6 +606,15 @@ version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb"
+[[package]]
+name = "foreign-types"
+version = "0.3.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1"
+dependencies = [
+ "foreign-types-shared 0.1.1",
+]
+
[[package]]
name = "foreign-types"
version = "0.5.0"
@@ -541,7 +622,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d737d9aa519fb7b749cbc3b962edcf310a8dd1f4b67c91c4f83975dbdd17d965"
dependencies = [
"foreign-types-macros",
- "foreign-types-shared",
+ "foreign-types-shared 0.3.1",
]
[[package]]
@@ -555,6 +636,12 @@ dependencies = [
"syn 3.0.6",
]
+[[package]]
+name = "foreign-types-shared"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b"
+
[[package]]
name = "foreign-types-shared"
version = "0.3.1"
@@ -763,10 +850,10 @@ version = "0.33.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7047791b5bc903b8cd963014b355f71dc9864a9a0b727057676c1dcae5cbc15"
dependencies = [
- "bitflags",
+ "bitflags 2.13.2",
"block",
"core-graphics-types",
- "foreign-types",
+ "foreign-types 0.5.0",
"log",
"objc",
"paste",
@@ -1016,6 +1103,12 @@ dependencies = [
"pest",
]
+[[package]]
+name = "pkg-config"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f6b464fbc74e149a392436b17d523f769e057cb6877f6a5c4618bc6f11800548"
+
[[package]]
name = "portable-atomic"
version = "1.15.0"
@@ -1187,13 +1280,29 @@ dependencies = [
"crossbeam-utils",
]
+[[package]]
+name = "rdev"
+version = "0.5.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "00552ca2dc2f93b84cd7b5581de49549411e4e41d89e1c691bcb93dc4be360c3"
+dependencies = [
+ "cocoa",
+ "core-foundation 0.7.0",
+ "core-foundation-sys 0.7.0",
+ "core-graphics 0.19.2",
+ "lazy_static",
+ "libc",
+ "winapi",
+ "x11",
+]
+
[[package]]
name = "redox_syscall"
version = "0.5.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
dependencies = [
- "bitflags",
+ "bitflags 2.13.2",
]
[[package]]
@@ -1222,7 +1331,7 @@ version = "1.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d"
dependencies = [
- "bitflags",
+ "bitflags 2.13.2",
"errno",
"libc",
"linux-raw-sys",
@@ -1755,6 +1864,22 @@ dependencies = [
"winapi-util",
]
+[[package]]
+name = "winapi"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
[[package]]
name = "winapi-util"
version = "0.1.11"
@@ -1764,6 +1889,12 @@ dependencies = [
"windows-sys",
]
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
+
[[package]]
name = "windows-link"
version = "0.2.1"
@@ -1788,6 +1919,16 @@ dependencies = [
"memchr",
]
+[[package]]
+name = "x11"
+version = "2.21.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "502da5464ccd04011667b11c435cb992822c2c0dbde1770c988480d312a0db2e"
+dependencies = [
+ "libc",
+ "pkg-config",
+]
+
[[package]]
name = "xattr"
version = "1.6.1"
diff --git a/Cargo.toml b/Cargo.toml
index b66765d..fe81910 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,3 +7,4 @@ tract = "0.23"
tract-metal = "0.23"
ndarray = "0.16"
anyhow = "1.0"
+rdev = "0.5.3"
\ No newline at end of file
diff --git a/backup/latency.rs b/backup/latency.rs
new file mode 100644
index 0000000..58af088
--- /dev/null
+++ b/backup/latency.rs
@@ -0,0 +1,2 @@
+// moved to backup during revert
+
diff --git a/backup/latency_compare.rs b/backup/latency_compare.rs
new file mode 100644
index 0000000..58af088
--- /dev/null
+++ b/backup/latency_compare.rs
@@ -0,0 +1,2 @@
+// moved to backup during revert
+
diff --git a/model/policy.onnx b/model/policy.onnx
index a40f576..683123d 100644
Binary files a/model/policy.onnx and b/model/policy.onnx differ
diff --git a/model/policy.onnx.data b/model/policy.onnx.data
index e1afb2e..aa8db2a 100644
Binary files a/model/policy.onnx.data and b/model/policy.onnx.data differ
diff --git a/model/policy.pt b/model/policy.pt
new file mode 100644
index 0000000..09d0566
Binary files /dev/null and b/model/policy.pt differ
diff --git a/src/bin/latency.rs b/src/bin/latency.rs
new file mode 100644
index 0000000..72c22f9
--- /dev/null
+++ b/src/bin/latency.rs
@@ -0,0 +1,35 @@
+// A small latency measurement binary that calls the `policy` function 1000 times
+// and prints average latency per call. This file includes the existing
+// `src/policy.rs` implementation so it can call the same function directly.
+
+include!("../policy.rs");
+
+use std::time::Instant;
+
+fn main() -> std::io::Result<()> {
+ // Dummy observation of length 17
+ let obs: [f32; 17] = [0.0f32; 17];
+
+ // Warm-up iterations (model loading might be heavy on first call)
+ println!("Warming up (10 iterations)...");
+ for _ in 0..10 {
+ policy(obs).expect("policy failed during warmup");
+ }
+
+ let iterations: usize = 1000;
+ println!("Running {} iterations...", iterations);
+ let start = Instant::now();
+ for _ in 0..iterations {
+ policy(obs).expect("policy failed during run");
+ }
+ let elapsed = start.elapsed();
+
+ let total_ms = elapsed.as_secs_f64() * 1000.0;
+ let avg_ms = total_ms / (iterations as f64);
+
+ println!("Total elapsed: {:.3} ms", total_ms);
+ println!("Average per call: {:.6} ms", avg_ms);
+
+ Ok(())
+}
+
diff --git a/src/bin/latency_compare.rs b/src/bin/latency_compare.rs
new file mode 100644
index 0000000..b4faba2
--- /dev/null
+++ b/src/bin/latency_compare.rs
@@ -0,0 +1,66 @@
+use std::time::Instant;
+
+// Reuse existing `policy` which currently reloads the model each call (full pipeline).
+include!("../policy.rs");
+
+fn main() -> Result<()> {
+ let obs: [f32; 17] = [0.0f32; 17];
+
+ // Warm-up: call the full-pipeline policy a few times
+ println!("Warm-up full pipeline (5)...");
+ for _ in 0..5 {
+ policy(obs)?;
+ }
+
+ let iterations = 1000usize;
+
+ // Measure full pipeline (what you said you wanted)
+ println!("Measuring full pipeline ({} iterations)...", iterations);
+ let start_full = Instant::now();
+ for _ in 0..iterations {
+ policy(obs)?;
+ }
+ let elapsed_full = start_full.elapsed();
+ let total_ms_full = elapsed_full.as_secs_f64() * 1000.0;
+ let avg_ms_full = total_ms_full / (iterations as f64);
+
+ println!("Full pipeline total: {:.3} ms", total_ms_full);
+ println!("Full pipeline avg: {:.6} ms", avg_ms_full);
+
+ // Now measure inference-only by loading the model once and running the runnable repeatedly
+ println!("\nPreparing model & runnable for inference-only measurement...");
+ let model_path = std::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()?;
+
+ // basic checks (similar to policy.rs)
+ ensure!(model.input_count()? == 1, "policy model must have exactly one input");
+
+ let runtime = tract::runtime_for_name("metal")?;
+ let mut runnable = runtime.prepare(model)?;
+
+ // Warm-up runnable
+ println!("Warm-up runnable (5)...");
+ for _ in 0..5 {
+ let input = Tensor::from_slice(&[1usize, 17usize], &obs)?;
+ let _outputs = runnable.run([input])?;
+ }
+
+ // Measure inference-only
+ println!("Measuring inference-only ({} iterations)...", iterations);
+ let start_inf = Instant::now();
+ for _ in 0..iterations {
+ let input = Tensor::from_slice(&[1usize, 17usize], &obs)?;
+ let _outputs = runnable.run([input])?;
+ }
+ let elapsed_inf = start_inf.elapsed();
+ let total_ms_inf = elapsed_inf.as_secs_f64() * 1000.0;
+ let avg_ms_inf = total_ms_inf / (iterations as f64);
+
+ println!("Inference-only total: {:.3} ms", total_ms_inf);
+ println!("Inference-only avg: {:.6} ms", avg_ms_inf);
+
+ Ok(())
+}
+
diff --git a/src/command.rs b/src/command.rs
new file mode 100644
index 0000000..6dfddb9
--- /dev/null
+++ b/src/command.rs
@@ -0,0 +1,96 @@
+use std::sync::{
+ Arc,
+ atomic::{AtomicU8, Ordering},
+};
+use std::thread::{self, JoinHandle};
+
+/// 使用 rdev 监听系统级键盘事件。
+/// bit0 = Up, bit1 = Left, bit2 = Right
+pub fn spawn_listener() -> (JoinHandle<()>, Arc) {
+ let state = Arc::new(AtomicU8::new(0));
+ let state_cloned = Arc::clone(&state);
+
+ let handle = thread::spawn(move || {
+ if let Err(err) = rdev::listen(move |event| {
+ match event.event_type {
+ rdev::EventType::KeyPress(key) => {
+ let bit = match key {
+ rdev::Key::UpArrow => Some(0x1),
+ rdev::Key::LeftArrow => Some(0x2),
+ rdev::Key::RightArrow => Some(0x4),
+ _ => None,
+ };
+
+ if let Some(mask) = bit {
+ state_cloned.fetch_or(mask, Ordering::SeqCst);
+ let s = state_cloned.load(Ordering::SeqCst);
+ println!("[CMD] {:?} (state=0b{:03b})", state_to_vector(s), s);
+ }
+ }
+ rdev::EventType::KeyRelease(key) => {
+ let bit = match key {
+ rdev::Key::UpArrow => Some(0x1),
+ rdev::Key::LeftArrow => Some(0x2),
+ rdev::Key::RightArrow => Some(0x4),
+ _ => None,
+ };
+
+ if let Some(mask) = bit {
+ state_cloned.fetch_and(!mask, Ordering::SeqCst);
+ let s = state_cloned.load(Ordering::SeqCst);
+ println!("[CMD] {:?} (state=0b{:03b})", state_to_vector(s), s);
+ }
+ }
+ _ => {}
+ }
+ }) {
+ eprintln!("rdev listen failed: {:?}", err);
+ }
+ });
+
+ (handle, state)
+}
+
+pub fn state_from_arc(state: &Arc) -> [f32; 3] {
+ state_to_vector(state.load(Ordering::SeqCst))
+}
+
+/// 状态位转成控制向量。
+pub fn state_to_vector(state: u8) -> [f32; 3] {
+ if state & 0x1 != 0 {
+ if state & 0x2 != 0 {
+ [0.2, 0.0, -0.5]
+ } else if state & 0x4 != 0 {
+ [0.2, 0.0, 0.5]
+ } else {
+ [0.2, 0.0, 0.0]
+ }
+ } else {
+ [0.0, 0.0, 0.0]
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_no_key() {
+ assert_eq!(state_to_vector(0), [0.0, 0.0, 0.0]);
+ }
+
+ #[test]
+ fn test_up_only() {
+ assert_eq!(state_to_vector(0x1), [0.2, 0.0, 0.0]);
+ }
+
+ #[test]
+ fn test_up_left() {
+ assert_eq!(state_to_vector(0x1 | 0x2), [0.2, 0.0, -0.5]);
+ }
+
+ #[test]
+ fn test_up_right() {
+ assert_eq!(state_to_vector(0x1 | 0x4), [0.2, 0.0, 0.5]);
+ }
+}
diff --git a/src/main.rs b/src/main.rs
index 041a8fa..fe170cf 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,35 +1,58 @@
mod policy;
mod pingEsp32;
+mod command;
-use anyhow::{Context, Result};
-fn main() -> Result<()> {
- //while循环
- let mut i = 0;
- while i < 10 {
- let action = policy::policy([i as f32 *0.05; 36])?;
- println!("模型输出: {action:?}");
- i = i + 1;
+use std::thread::sleep;
+use pingEsp32::Link;
+use std::time::Duration;
+
+fn main() -> std::io::Result<()> {
+
+ let mut link = Link::open_default()?;
+ link.send_angles(&[1.57; 8])?;
+ //sleep10S
+ sleep(Duration::from_secs(10));
+
+ let (_handle, cmd_state) = command::spawn_listener();
+
+ loop {
+ // 阻塞等待下一条 ESP32 消息;收到后只处理一次 telemetry,并在处理完后再继续等待下一条
+ let msg = link.recv_timeout(Duration::from_millis(1000));
+ //接收imu数据并添加到观测中(obs)
+ if let Some(msg) = msg {
+ if msg.is_data() {
+ if let Some(obs) = msg.telemetry() {
+ //计时
+ // let start = std::time::Instant::now();
+ //添加动作指令到观测中(obs)
+ let cmd_vec = command::state_from_arc(&cmd_state);
+ let mut obs = obs;
+ obs[6] = cmd_vec[0];
+ obs[7] = cmd_vec[1];
+ obs[8] = cmd_vec[2];
+ //添加上一个步骤的动作到观测值中(obs)
+ let action = policy::policy(obs).expect("policy inference failed");
+ link.send_angles(&action)?;
+ println!("[Action]: {:?}", action);
+ obs[9] = action[0];
+ obs[10] = action[1];
+ obs[11] = action[2];
+ obs[12] = action[3];
+ obs[13] = action[4];
+ obs[14] = action[5];
+ obs[15] = action[6];
+ obs[16] = action[7];
+ println!("发送 17 维观测值(含动作): {:?}", obs);
+ //结束计时
+ // let duration = start.elapsed();
+ // println!("程序运行时间: {:?}", duration);
+ continue;
+ }
+ } else if msg.is_reply() {
+ println!("[REPLY] {}", msg.text());
+ } else {
+ println!("[TEXT] {}", msg.text());
+ }
+ }
}
- Ok(())
}
-
-use std::time::Instant;
-
-fn main1() -> Result<()> {
- const N: u32 = 500;
- 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(())
-}
\ No newline at end of file
diff --git a/src/pingEsp32.rs b/src/pingEsp32.rs
index e69de29..408e717 100644
--- a/src/pingEsp32.rs
+++ b/src/pingEsp32.rs
@@ -0,0 +1,281 @@
+use std::io::{BufRead, BufReader, Read, Write};
+use std::net::TcpStream;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::mpsc::{self, Receiver, Sender, RecvTimeoutError};
+use std::sync::Arc;
+use std::thread::{self, JoinHandle};
+use std::time::Duration;
+
+use std::f32::consts::PI;
+
+pub const ESP32_IP: &str = "192.168.4.1";
+pub const ESP32_PORT: u16 = 3333;
+pub const NUM_SERVOS: usize = 8;
+
+const TELEMETRY_MAGIC: [u8; 2] = [0xA5, 0x5A];
+const TELEMETRY_FRAME_LEN: usize = 16;
+
+/// 从 ESP32 收到的一条消息(可包含 ASCII 文本,也可包含 36 维 IMU 数据)
+#[derive(Debug, Clone)]
+pub struct Esp32Message {
+ text: String,
+ telemetry: Option<[f32; 17]>,
+}
+
+impl Esp32Message {
+ pub fn new_text(text: impl Into) -> Self {
+ Self {
+ text: text.into(),
+ telemetry: None,
+ }
+ }
+
+ pub fn new_telemetry(data: [f32; 17]) -> Self {
+ Self {
+ text: "IMU".to_string(),
+ telemetry: Some(data),
+ }
+ }
+
+ /// 是否是数据推送(二进制 IMU 帧)
+ pub fn is_data(&self) -> bool { self.telemetry.is_some() }
+
+ /// 是否是应答,比如 "ESP32: OK ..."
+ pub fn is_reply(&self) -> bool {
+ self.telemetry.is_none() && self.text.starts_with("ESP32:")
+ }
+
+ pub fn text(&self) -> &str { &self.text }
+
+ /// 返回 36 维 f32 数组;其余位置按 0.0 填充。
+ pub fn telemetry(&self) -> Option<[f32; 17]> { self.telemetry }
+}
+
+// ==================== 双向连接 ====================
+
+pub struct Link {
+ writer: TcpStream,
+ rx: Receiver,
+ stop: Arc,
+ reader_job: Option>,
+}
+
+impl Link {
+ /// 建立连接并启动后台接收线程
+ pub fn open(ip: &str, port: u16) -> std::io::Result {
+ let addr = format!("{}:{}", ip, port);
+ println!("[*] 连接 {} ...", addr);
+ println!("请连接上esp32的wifi热点后再运行本程序【ESP32_AP】");
+ let stream = TcpStream::connect(&addr)?;
+ stream.set_nodelay(true)?;
+ stream.set_read_timeout(Some(Duration::from_millis(200)))?;
+ println!("[+] 已连接");
+
+ let reader_stream = stream.try_clone()?;
+ let (tx, rx): (Sender, Receiver) = mpsc::channel();
+ let stop = Arc::new(AtomicBool::new(false));
+
+ let reader_job = spawn_reader(reader_stream, tx, Arc::clone(&stop));
+
+ Ok(Self {
+ writer: stream,
+ rx,
+ stop,
+ reader_job: Some(reader_job),
+ })
+ }
+
+ pub fn open_default() -> std::io::Result {
+ Self::open(ESP32_IP, ESP32_PORT)
+ }
+
+ // ---------- 发送 ----------
+
+ /// 发送弧度列表
+ pub fn send_angles(&mut self, rads: &[f32]) -> std::io::Result<()> {
+ let n = rads.len().min(NUM_SERVOS);
+ let payload = rads[..n]
+ .iter()
+ .map(|r| format!("{:.4}", r.clamp(0.0, PI)))
+ .collect::>()
+ .join(",");
+ self.send_line(&format!("RAD:{}", payload))
+ }
+
+ /// 发送角度列表
+ pub fn send_degrees(&mut self, degs: &[f32]) -> std::io::Result<()> {
+ let n = degs.len().min(NUM_SERVOS);
+ let payload = degs[..n]
+ .iter()
+ .map(|d| format!("{:.2}", d.clamp(0.0, 180.0)))
+ .collect::>()
+ .join(",");
+ self.send_line(&format!("DEG:{}", payload))
+ }
+
+ /// 发送任意一行(自动补 '\n')
+ pub fn send_line(&mut self, line: &str) -> std::io::Result<()> {
+ self.writer.write_all(line.as_bytes())?;
+ self.writer.write_all(b"\n")?;
+ self.writer.flush()
+ }
+
+ // ---------- 接收 ----------
+
+ /// 阻塞等待一条消息(最多等 timeout)
+ pub fn recv_timeout(&self, timeout: Duration) -> Option {
+ match self.rx.recv_timeout(timeout) {
+ Ok(m) => Some(m),
+ Err(RecvTimeoutError::Timeout) => None,
+ Err(RecvTimeoutError::Disconnected) => None,
+ }
+ }
+
+ /// 非阻塞:立刻返回当前队列里有没有消息
+ pub fn try_recv(&self) -> Option {
+ self.rx.try_recv().ok()
+ }
+}
+
+impl Drop for Link {
+ fn drop(&mut self) {
+ self.stop.store(true, Ordering::SeqCst);
+ let _ = self.writer.shutdown(std::net::Shutdown::Both);
+ if let Some(job) = self.reader_job.take() {
+ let _ = job.join();
+ }
+ }
+}
+
+// ==================== 后台接收线程 ====================
+
+fn spawn_reader(
+ stream: TcpStream,
+ tx: Sender,
+ stop: Arc,
+) -> JoinHandle<()> {
+ thread::spawn(move || {
+ let mut stream = stream;
+ let mut pending = Vec::new();
+ let mut buffer = [0u8; 1024];
+
+ loop {
+ if stop.load(Ordering::SeqCst) {
+ break;
+ }
+
+ match stream.read(&mut buffer) {
+ Ok(0) => break,
+ Ok(n) => {
+ pending.extend_from_slice(&buffer[..n]);
+
+ loop {
+ if pending.len() >= TELEMETRY_FRAME_LEN
+ && pending[..2] == TELEMETRY_MAGIC
+ {
+ let frame = &pending[..TELEMETRY_FRAME_LEN];
+ if let Some(data) = parse_telemetry_frame(frame) {
+ if tx.send(Esp32Message::new_telemetry(data)).is_err() {
+ return;
+ }
+ pending.drain(..TELEMETRY_FRAME_LEN);
+ continue;
+ }
+
+ // CRC 错误,丢掉首字节继续找下一帧
+ pending.drain(..1);
+ continue;
+ }
+
+ if let Some(newline_pos) = pending.iter().position(|&b| b == b'\n') {
+ let line_bytes = &pending[..newline_pos];
+ // 只处理纯 ASCII 文本; 非法 UTF-8 也不阻断连接
+ let line = String::from_utf8_lossy(line_bytes).trim().to_string();
+ if !line.is_empty() {
+ if tx.send(Esp32Message::new_text(line)).is_err() {
+ return;
+ }
+ }
+ pending.drain(..=newline_pos);
+ continue;
+ }
+
+ // 还没凑齐完整帧/整行,等下一批数据
+ break;
+ }
+
+ if pending.len() > 65536 {
+ pending.clear();
+ }
+ }
+ Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock || e.kind() == std::io::ErrorKind::TimedOut => {
+ continue;
+ }
+ Err(_) => break,
+ }
+ }
+
+ println!("[*] 接收线程退出");
+ })
+}
+
+fn parse_telemetry_frame(frame: &[u8]) -> Option<[f32; 17]> {
+ // Frame layout (from ESP32):
+ // 0..1 magic 0xA5 0x5A
+ // 2..3 gravity.x (int16 big-endian) = projection * 10000
+ // 4..5 gravity.y
+ // 6..7 gravity.z
+ // 8..9 gyro_x encoded (int16 BE) = rad/s * 1000
+ // 10..11 gyro_y
+ // 12..13 gyro_z
+ // 14..15 CRC16 over bytes 0..13 (big-endian)
+ if frame.len() != TELEMETRY_FRAME_LEN {
+ return None;
+ }
+ if frame[..2] != TELEMETRY_MAGIC {
+ return None;
+ }
+
+ let expected_crc = u16::from_be_bytes([frame[14], frame[15]]);
+ let actual_crc = calculate_crc16_ccitt(&frame[..14]);
+ if actual_crc != expected_crc {
+ return None;
+ }
+
+ // Decode gravity projections (stored as projection * 10000)
+ let gx_proj = i16::from_be_bytes([frame[2], frame[3]]) as f32 / 10000.0;
+ let gy_proj = i16::from_be_bytes([frame[4], frame[5]]) as f32 / 10000.0;
+ let gz_proj = i16::from_be_bytes([frame[6], frame[7]]) as f32 / 10000.0;
+
+ // Decode base angular velocities (stored as rad/s * 1000)
+ let avx = i16::from_be_bytes([frame[8], frame[9]]) as f32 / 1000.0;
+ let avy = i16::from_be_bytes([frame[10], frame[11]]) as f32 / 1000.0;
+ let avz = i16::from_be_bytes([frame[12], frame[13]]) as f32 / 1000.0;
+
+ let mut values = [0.0f32; 17];
+ // fill first 6 entries: gravity x,y,z then angular vel x,y,z
+ values[0] = gx_proj;
+ values[1] = gy_proj;
+ values[2] = gz_proj;
+ values[3] = avx;
+ values[4] = avy;
+ values[5] = avz;
+
+ Some(values)
+}
+
+fn calculate_crc16_ccitt(data: &[u8]) -> u16 {
+ let mut crc = 0xFFFFu16;
+ for &byte in data {
+ crc ^= (byte as u16) << 8;
+ for _ in 0..8 {
+ if crc & 0x8000 != 0 {
+ crc = ((crc << 1) ^ 0x1021) & 0xFFFF;
+ } else {
+ crc = (crc << 1) & 0xFFFF;
+ }
+ }
+ }
+ crc
+}
+
diff --git a/src/policy.rs b/src/policy.rs
index 729754e..32142d6 100644
--- a/src/policy.rs
+++ b/src/policy.rs
@@ -7,7 +7,7 @@ use tract::prelude::*;
/// - 每次调用都会重新加载并准备模型(简单、无状态,方便一行调用)。
/// - 模型路径固定为 `CARGO_MANIFEST_DIR/model/policy.onnx`。
/// - 使用 Metal 运行时。
-pub fn policy(observation: [f32; 36]) -> Result<[f32; 8]> {
+pub fn policy(observation: [f32; 17]) -> Result<[f32; 8]> {
let model_path = Path::new(env!("CARGO_MANIFEST_DIR")).join("model/policy.onnx");
// 加载模型
@@ -66,6 +66,6 @@ pub fn policy(observation: [f32; 36]) -> Result<[f32; 8]> {
result.copy_from_slice(values);
result
.iter_mut()
- .for_each(|value| *value = (*value * 0.5).clamp(-1.57, 1.57));
+ .for_each(|value| *value = (*value * 0.5 + 1.57).clamp(0.0, 3.14));
Ok(result)
}
\ No newline at end of file