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