diff --git a/Cargo.lock b/Cargo.lock index 0748530c4368..3cad1a73d975 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3110,7 +3110,7 @@ dependencies = [ [[package]] name = "dimos-viewer" -version = "0.30.0-alpha.1+dev" +version = "0.30.0-alpha.4" dependencies = [ "bincode", "clap", diff --git a/dimos/src/interaction/keyboard.rs b/dimos/src/interaction/keyboard.rs index b6cdcd809c37..26792202f6d5 100644 --- a/dimos/src/interaction/keyboard.rs +++ b/dimos/src/interaction/keyboard.rs @@ -1,247 +1,209 @@ //! Keyboard handler for WASD movement controls that publish Twist messages. -//! -//! Converts keyboard input to robot velocity commands following teleop conventions: -//! - WASD/arrows for linear/angular motion -//! - QE for strafing -//! - Space for emergency stop -//! - Shift for speed multiplier +//! +//! Supports multi-robot: tracks which robot is "active" for teleop. +//! WASD publishes to that robot's /cmd_vel channel. use std::io; use super::lcm::{LcmPublisher, twist_command}; use rerun::external::{egui, re_log}; -/// LCM channel for twist commands (matches DimOS convention) -const CMD_VEL_CHANNEL: &str = "/cmd_vel#geometry_msgs.Twist"; +/// LCM channel suffix for twist commands +const CMD_VEL_SUFFIX: &str = "/cmd_vel#geometry_msgs.Twist"; +/// Default channel when no robot is selected +const DEFAULT_CMD_VEL: &str = "/cmd_vel#geometry_msgs.Twist"; -/// Base speeds for keyboard control -const BASE_LINEAR_SPEED: f64 = 0.5; // m/s -const BASE_ANGULAR_SPEED: f64 = 0.8; // rad/s -const FAST_MULTIPLIER: f64 = 2.0; // Shift modifier +/// Base speeds +const BASE_LINEAR_SPEED: f64 = 0.5; +const BASE_ANGULAR_SPEED: f64 = 0.8; +const FAST_MULTIPLIER: f64 = 2.0; /// Overlay styling -const OVERLAY_MARGIN: f32 = 12.0; +const OVERLAY_X: f32 = 260.0; +const OVERLAY_Y: f32 = 12.0; const OVERLAY_PADDING: f32 = 10.0; const OVERLAY_ROUNDING: f32 = 8.0; -const OVERLAY_BG: egui::Color32 = egui::Color32::from_rgba_premultiplied(20, 20, 30, 220); +const OVERLAY_BG_ACTIVE: egui::Color32 = egui::Color32::from_rgba_premultiplied(20, 20, 30, 220); +const OVERLAY_BG_IDLE: egui::Color32 = egui::Color32::from_rgba_premultiplied(20, 20, 30, 100); const KEY_SIZE: f32 = 32.0; const KEY_GAP: f32 = 3.0; const KEY_ACTIVE_BG: egui::Color32 = egui::Color32::from_rgb(60, 180, 75); const KEY_INACTIVE_BG: egui::Color32 = egui::Color32::from_rgba_premultiplied(60, 60, 80, 180); +const KEY_IDLE_BG: egui::Color32 = egui::Color32::from_rgba_premultiplied(60, 60, 80, 80); const KEY_TEXT_COLOR: egui::Color32 = egui::Color32::WHITE; +const KEY_TEXT_IDLE: egui::Color32 = egui::Color32::from_rgba_premultiplied(255, 255, 255, 100); const LABEL_COLOR: egui::Color32 = egui::Color32::from_rgb(180, 180, 200); +const LABEL_IDLE: egui::Color32 = egui::Color32::from_rgba_premultiplied(180, 180, 200, 80); const ESTOP_ACTIVE_BG: egui::Color32 = egui::Color32::from_rgb(220, 50, 50); -/// Tracks which movement keys are currently held down. #[derive(Debug, Clone, Default)] struct KeyState { - forward: bool, // W or Up - backward: bool, // S or Down - left: bool, // A or Left - right: bool, // D or Right - strafe_l: bool, // Q - strafe_r: bool, // E - fast: bool, // Shift held + forward: bool, + backward: bool, + left: bool, + right: bool, + strafe_l: bool, + strafe_r: bool, + fast: bool, } impl KeyState { - fn new() -> Self { - Default::default() - } + fn new() -> Self { Default::default() } - /// Returns true if any movement key is currently active fn any_active(&self) -> bool { self.forward || self.backward || self.left || self.right || self.strafe_l || self.strafe_r } - /// Reset all key states (used for emergency stop) - fn reset(&mut self) { - self.forward = false; - self.backward = false; - self.left = false; - self.right = false; - self.strafe_l = false; - self.strafe_r = false; - self.fast = false; - } + fn reset(&mut self) { *self = Default::default(); } } /// Handles keyboard input and publishes Twist via LCM. +/// Supports multi-robot via `set_active_robot()`. pub struct KeyboardHandler { publisher: LcmPublisher, state: KeyState, was_active: bool, - estop_flash: bool, // true briefly after space pressed + estop_flash: bool, + /// Whether teleop is engaged (robot clicked in 3D view) + engaged: bool, + /// Currently selected robot entity path prefix, or None for default /cmd_vel + active_robot: Option, } impl KeyboardHandler { - /// Create a new keyboard handler with LCM publisher on CMD_VEL_CHANNEL. pub fn new() -> Result { - let publisher = LcmPublisher::new(CMD_VEL_CHANNEL.to_string())?; + let publisher = LcmPublisher::new(DEFAULT_CMD_VEL.to_string())?; Ok(Self { publisher, state: KeyState::new(), was_active: false, estop_flash: false, + engaged: false, + active_robot: None, }) } - /// Process keyboard input from egui and publish Twist if keys are held. - /// Called once per frame from DimosApp.ui(). - /// - /// Returns true if any movement key is active (for UI overlay). + /// Set the active robot for teleop. Recreates LCM publisher for new channel. + pub fn set_active_robot(&mut self, robot_prefix: Option) { + if self.active_robot == robot_prefix { return; } + if self.was_active { + let _ = self.publish_stop(); + self.was_active = false; + } + let channel = match &robot_prefix { + Some(prefix) => format!("{prefix}{CMD_VEL_SUFFIX}"), + None => DEFAULT_CMD_VEL.to_string(), + }; + match LcmPublisher::new(channel.clone()) { + Ok(p) => { self.publisher = p; re_log::info!("Teleop target: {channel}"); } + Err(e) => { re_log::error!("Publisher failed for {channel}: {e}"); return; } + } + self.active_robot = robot_prefix; + } + + pub fn active_robot(&self) -> Option<&str> { self.active_robot.as_deref() } + + /// Whether teleop is currently engaged (robot selected). + pub fn engaged(&self) -> bool { self.engaged } + + /// Set engaged state. Sends stop when disengaging. + pub fn set_engaged(&mut self, engaged: bool) { + if self.engaged && !engaged { + let _ = self.publish_stop(); + self.state.reset(); + self.was_active = false; + } + self.engaged = engaged; + } + pub fn process(&mut self, ctx: &egui::Context) -> bool { self.estop_flash = false; - - // Check if any text widget has focus - if so, skip keyboard capture + // Only capture keys when engaged and no text field focused let text_has_focus = ctx.memory(|m| m.focused().is_some()); - if text_has_focus { - if self.was_active { - if let Err(e) = self.publish_stop() { - re_log::warn!("Failed to send stop command on focus change: {e:?}"); - } - self.was_active = false; - } + if !self.engaged || text_has_focus { + if self.was_active { let _ = self.publish_stop(); self.was_active = false; } return false; } - - // Update key state from egui input self.update_key_state(ctx); - - // Check for emergency stop (Space key pressed - one-shot action) if ctx.input(|i| i.key_pressed(egui::Key::Space)) { self.state.reset(); - if let Err(e) = self.publish_stop() { - re_log::warn!("Failed to send emergency stop: {e:?}"); - } + let _ = self.publish_stop(); self.was_active = false; self.estop_flash = true; - return true; // return true so overlay shows the e-stop flash + return true; } - - // Publish twist command if keys are active, or stop if just released if self.state.any_active() { - if let Err(e) = self.publish_twist() { - re_log::warn!("Failed to publish twist command: {e:?}"); - } + let _ = self.publish_twist(); self.was_active = true; } else if self.was_active { - if let Err(e) = self.publish_stop() { - re_log::warn!("Failed to send stop on key release: {e:?}"); - } + let _ = self.publish_stop(); self.was_active = false; } - self.state.any_active() } - /// Draw keyboard overlay HUD. Always shown (dim when idle, bright when active). + /// Draw keyboard overlay HUD. Greyed out when idle, bright when active. pub fn draw_overlay(&self, ctx: &egui::Context) { egui::Area::new("keyboard_hud".into()) - .fixed_pos(egui::pos2(OVERLAY_MARGIN, OVERLAY_MARGIN)) + .default_pos(egui::pos2(OVERLAY_X, OVERLAY_Y)) + .movable(true) .order(egui::Order::Foreground) .interactable(false) .show(ctx, |ui| { + let bg = if self.engaged { OVERLAY_BG_ACTIVE } else { OVERLAY_BG_IDLE }; egui::Frame::new() - .fill(OVERLAY_BG) + .fill(bg) .corner_radius(egui::CornerRadius::same(OVERLAY_ROUNDING as u8)) .inner_margin(egui::Margin::same(OVERLAY_PADDING as i8)) - .show(ui, |ui| { - self.draw_hud_content(ui); - }); + .show(ui, |ui| self.draw_hud_content(ui)); }); } fn draw_hud_content(&self, ui: &mut egui::Ui) { - let active = self.state.any_active() || self.estop_flash; - - // Title - let title_color = if active { - egui::Color32::WHITE + let (title_color, label_color) = if self.engaged { + (egui::Color32::WHITE, LABEL_COLOR) } else { - egui::Color32::from_rgb(120, 120, 140) + (egui::Color32::from_rgba_premultiplied(255, 255, 255, 80), LABEL_IDLE) }; - ui.label(egui::RichText::new("🎮 Keyboard Teleop").color(title_color).size(13.0)); - ui.add_space(4.0); - - // Key grid: [Q] [W] [E] - // [A] [S] [D] - // [ SPACE ] - let row1 = [ - ("Q", self.state.strafe_l), - ("W", self.state.forward), - ("E", self.state.strafe_r), - ]; - let row2 = [ - ("A", self.state.left), - ("S", self.state.backward), - ("D", self.state.right), - ]; - - // Row 1 - ui.horizontal(|ui| { - ui.spacing_mut().item_spacing.x = KEY_GAP; - for (label, pressed) in &row1 { - self.draw_key(ui, label, *pressed); - } - }); - - // Row 2 - ui.horizontal(|ui| { - ui.spacing_mut().item_spacing.x = KEY_GAP; - for (label, pressed) in &row2 { - self.draw_key(ui, label, *pressed); + let title = if self.engaged { + match &self.active_robot { + Some(name) => format!("🎮 {}", name.rsplit('/').next().unwrap_or(name)), + None => "🎮 Teleop".to_string(), } - }); - - // Space bar (e-stop) - let space_width = KEY_SIZE * 3.0 + KEY_GAP * 2.0; - let space_rect = ui.allocate_exact_size( - egui::vec2(space_width, KEY_SIZE * 0.7), - egui::Sense::hover(), - ).0; - let space_bg = if self.estop_flash { - ESTOP_ACTIVE_BG } else { - KEY_INACTIVE_BG + "🎮 Teleop (click robot)".to_string() }; - ui.painter().rect_filled(space_rect, egui::CornerRadius::same(4), space_bg); - ui.painter().text( - space_rect.center(), - egui::Align2::CENTER_CENTER, - "STOP", - egui::FontId::proportional(11.0), - KEY_TEXT_COLOR, - ); - + ui.label(egui::RichText::new(title).color(title_color).size(13.0)); ui.add_space(4.0); - // Speed indicator - let speed_label = if self.state.fast { "⇧ FAST" } else { "⇧ shift=fast" }; - let speed_color = if self.state.fast { - egui::Color32::from_rgb(255, 200, 50) - } else { - LABEL_COLOR - }; + let key_bg = if self.engaged { KEY_INACTIVE_BG } else { KEY_IDLE_BG }; + let text_col = if self.engaged { KEY_TEXT_COLOR } else { KEY_TEXT_IDLE }; + let rows = [ + [("Q", self.state.strafe_l), ("W", self.state.forward), ("E", self.state.strafe_r)], + [("A", self.state.left), ("S", self.state.backward), ("D", self.state.right)], + ]; + for row in &rows { + ui.horizontal(|ui| { + ui.spacing_mut().item_spacing.x = KEY_GAP; + for (label, pressed) in row { + let (rect, _) = ui.allocate_exact_size(egui::vec2(KEY_SIZE, KEY_SIZE), egui::Sense::hover()); + let bg = if *pressed { KEY_ACTIVE_BG } else { key_bg }; + ui.painter().rect_filled(rect, egui::CornerRadius::same(4), bg); + ui.painter().text(rect.center(), egui::Align2::CENTER_CENTER, label, egui::FontId::monospace(14.0), text_col); + } + }); + } + let space_w = KEY_SIZE * 3.0 + KEY_GAP * 2.0; + let (space_rect, _) = ui.allocate_exact_size(egui::vec2(space_w, KEY_SIZE * 0.7), egui::Sense::hover()); + let space_bg = if self.estop_flash { ESTOP_ACTIVE_BG } else { key_bg }; + ui.painter().rect_filled(space_rect, egui::CornerRadius::same(4), space_bg); + ui.painter().text(space_rect.center(), egui::Align2::CENTER_CENTER, "STOP", egui::FontId::proportional(11.0), text_col); + ui.add_space(4.0); + let (speed_label, speed_color) = if self.state.fast { + ("⇧ FAST", egui::Color32::from_rgb(255, 200, 50)) + } else { ("⇧ shift=fast", label_color) }; ui.label(egui::RichText::new(speed_label).color(speed_color).size(10.0)); } - fn draw_key(&self, ui: &mut egui::Ui, label: &str, pressed: bool) { - let (rect, _) = ui.allocate_exact_size( - egui::vec2(KEY_SIZE, KEY_SIZE), - egui::Sense::hover(), - ); - let bg = if pressed { KEY_ACTIVE_BG } else { KEY_INACTIVE_BG }; - ui.painter().rect_filled(rect, egui::CornerRadius::same(4), bg); - ui.painter().text( - rect.center(), - egui::Align2::CENTER_CENTER, - label, - egui::FontId::monospace(14.0), - KEY_TEXT_COLOR, - ); - } - - /// Read current key state from egui input, update self.state. fn update_key_state(&mut self, ctx: &egui::Context) { ctx.input(|i| { self.state.forward = i.key_down(egui::Key::W) || i.key_down(egui::Key::ArrowUp); @@ -254,70 +216,34 @@ impl KeyboardHandler { }); } - /// Convert current KeyState to Twist and publish via LCM. fn publish_twist(&mut self) -> io::Result<()> { - let (lin_x, lin_y, lin_z, ang_x, ang_y, ang_z) = self.compute_twist(); - - let cmd = twist_command( - [lin_x, lin_y, lin_z], - [ang_x, ang_y, ang_z], - ); - - self.publisher.publish_twist(&cmd)?; - - re_log::trace!( - "Published twist: lin=({:.2},{:.2},{:.2}) ang=({:.2},{:.2},{:.2})", - lin_x, lin_y, lin_z, ang_x, ang_y, ang_z - ); - - Ok(()) + let (lx, ly, lz, ax, ay, az) = self.compute_twist(); + self.publisher.publish_twist(&twist_command([lx, ly, lz], [ax, ay, az])).map(|_| ()) } - /// Publish all-zero twist (stop command) fn publish_stop(&mut self) -> io::Result<()> { - let cmd = twist_command([0.0, 0.0, 0.0], [0.0, 0.0, 0.0]); - self.publisher.publish_twist(&cmd)?; - re_log::debug!("Published stop command"); - Ok(()) + self.publisher.publish_twist(&twist_command([0.0; 3], [0.0; 3])).map(|_| ()) } - /// Map KeyState to linear/angular velocities. fn compute_twist(&self) -> (f64, f64, f64, f64, f64, f64) { - let mut linear_x = 0.0; - let mut linear_y = 0.0; - let mut angular_z = 0.0; - - if self.state.forward { - linear_x += BASE_LINEAR_SPEED; - } - if self.state.backward { - linear_x -= BASE_LINEAR_SPEED; - } - if self.state.strafe_l { - linear_y += BASE_LINEAR_SPEED; - } - if self.state.strafe_r { - linear_y -= BASE_LINEAR_SPEED; - } - if self.state.left { - angular_z += BASE_ANGULAR_SPEED; - } - if self.state.right { - angular_z -= BASE_ANGULAR_SPEED; - } - if self.state.fast { - linear_x *= FAST_MULTIPLIER; - linear_y *= FAST_MULTIPLIER; - angular_z *= FAST_MULTIPLIER; - } - - (linear_x, linear_y, 0.0, 0.0, 0.0, angular_z) + let mut lx = 0.0; + let mut ly = 0.0; + let mut az = 0.0; + if self.state.forward { lx += BASE_LINEAR_SPEED; } + if self.state.backward { lx -= BASE_LINEAR_SPEED; } + if self.state.strafe_l { ly += BASE_LINEAR_SPEED; } + if self.state.strafe_r { ly -= BASE_LINEAR_SPEED; } + if self.state.left { az += BASE_ANGULAR_SPEED; } + if self.state.right { az -= BASE_ANGULAR_SPEED; } + if self.state.fast { lx *= FAST_MULTIPLIER; ly *= FAST_MULTIPLIER; az *= FAST_MULTIPLIER; } + (lx, ly, 0.0, 0.0, 0.0, az) } } impl std::fmt::Debug for KeyboardHandler { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("KeyboardHandler") + .field("active_robot", &self.active_robot) .field("state", &self.state) .field("was_active", &self.was_active) .finish() @@ -332,179 +258,122 @@ mod tests { fn test_key_state_any_active() { let mut state = KeyState::new(); assert!(!state.any_active()); - state.forward = true; assert!(state.any_active()); - state.reset(); assert!(!state.any_active()); - - state.strafe_l = true; - assert!(state.any_active()); } #[test] fn test_wasd_to_twist_mapping() { - let mut state = KeyState::new(); - state.forward = true; - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { forward: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, BASE_LINEAR_SPEED); - assert_eq!(lin_y, 0.0); - assert_eq!(ang_z, 0.0); + let (lx, ly, _, _, _, az) = h.compute_twist(); + assert_eq!(lx, BASE_LINEAR_SPEED); + assert_eq!(ly, 0.0); + assert_eq!(az, 0.0); } #[test] fn test_turn_left_right_mapping() { - let mut state = KeyState::new(); - state.left = true; - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { left: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, 0.0); - assert_eq!(lin_y, 0.0); - assert_eq!(ang_z, BASE_ANGULAR_SPEED); - - let mut state = KeyState::new(); - state.right = true; - let handler = KeyboardHandler { + assert_eq!(h.compute_twist().5, BASE_ANGULAR_SPEED); + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { right: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, 0.0); - assert_eq!(lin_y, 0.0); - assert_eq!(ang_z, -BASE_ANGULAR_SPEED); + assert_eq!(h.compute_twist().5, -BASE_ANGULAR_SPEED); } #[test] fn test_strafe_mapping() { - let mut state = KeyState::new(); - state.strafe_l = true; - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { strafe_l: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, 0.0); - assert_eq!(lin_y, BASE_LINEAR_SPEED); - assert_eq!(ang_z, 0.0); - - let mut state = KeyState::new(); - state.strafe_r = true; - let handler = KeyboardHandler { - publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, - }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, 0.0); - assert_eq!(lin_y, -BASE_LINEAR_SPEED); - assert_eq!(ang_z, 0.0); + assert_eq!(h.compute_twist().1, BASE_LINEAR_SPEED); } #[test] fn test_shift_doubles_speed() { - let mut state = KeyState::new(); - state.forward = true; - state.fast = true; - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { forward: true, fast: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, BASE_LINEAR_SPEED * FAST_MULTIPLIER); - assert_eq!(lin_y, 0.0); - assert_eq!(ang_z, 0.0); + assert_eq!(h.compute_twist().0, BASE_LINEAR_SPEED * FAST_MULTIPLIER); } #[test] fn test_simultaneous_keys() { - let mut state = KeyState::new(); - state.forward = true; - state.left = true; - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { forward: true, left: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, BASE_LINEAR_SPEED); - assert_eq!(lin_y, 0.0); - assert_eq!(ang_z, BASE_ANGULAR_SPEED); + let t = h.compute_twist(); + assert_eq!(t.0, BASE_LINEAR_SPEED); + assert_eq!(t.5, BASE_ANGULAR_SPEED); } #[test] fn test_key_reset() { - let mut state = KeyState::new(); - state.forward = true; - state.left = true; - state.fast = true; - assert!(state.any_active()); - state.reset(); - assert!(!state.forward); - assert!(!state.left); - assert!(!state.fast); - assert!(!state.any_active()); + let mut s = KeyState { forward: true, left: true, fast: true, ..Default::default() }; + s.reset(); + assert!(!s.any_active()); + assert!(!s.fast); } #[test] fn test_keyboard_handler_creation() { - let handler = KeyboardHandler::new(); - assert!(handler.is_ok()); - let handler = handler.unwrap(); - assert!(!handler.was_active); - assert!(!handler.state.any_active()); + let h = KeyboardHandler::new().unwrap(); + assert!(!h.was_active); + assert!(!h.engaged); + assert!(h.active_robot.is_none()); } #[test] fn test_opposite_keys_cancel() { - let mut state = KeyState::new(); - state.forward = true; - state.backward = true; - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), - state, - was_active: false, - estop_flash: false, + state: KeyState { forward: true, backward: true, ..Default::default() }, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, _, _, _, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, 0.0); - assert_eq!(lin_y, 0.0); - assert_eq!(ang_z, 0.0); + assert_eq!(h.compute_twist().0, 0.0); } #[test] fn test_compute_twist_all_zeros() { - let handler = KeyboardHandler { + let h = KeyboardHandler { publisher: LcmPublisher::new("/test".to_string()).unwrap(), state: KeyState::new(), - was_active: false, - estop_flash: false, + was_active: false, estop_flash: false, engaged: true, active_robot: None, }; - let (lin_x, lin_y, lin_z, ang_x, ang_y, ang_z) = handler.compute_twist(); - assert_eq!(lin_x, 0.0); - assert_eq!(lin_y, 0.0); - assert_eq!(lin_z, 0.0); - assert_eq!(ang_x, 0.0); - assert_eq!(ang_y, 0.0); - assert_eq!(ang_z, 0.0); + let (lx, ly, lz, ax, ay, az) = h.compute_twist(); + assert!(lx == 0.0 && ly == 0.0 && lz == 0.0 && ax == 0.0 && ay == 0.0 && az == 0.0); + } + + #[test] + fn test_set_active_robot() { + let mut h = KeyboardHandler::new().unwrap(); + assert!(h.active_robot().is_none()); + assert!(!h.engaged()); + h.set_active_robot(Some("/world/go2".to_string())); + h.set_engaged(true); + assert_eq!(h.active_robot(), Some("/world/go2")); + assert!(h.engaged()); + h.set_engaged(false); + assert!(!h.engaged()); + h.set_active_robot(None); + assert!(h.active_robot().is_none()); } } diff --git a/dimos/src/viewer.rs b/dimos/src/viewer.rs index 7af7282ef188..f942740e81cc 100644 --- a/dimos/src/viewer.rs +++ b/dimos/src/viewer.rs @@ -1,4 +1,4 @@ -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; use std::rc::Rc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; @@ -10,256 +10,169 @@ use rerun::external::{eframe, egui, re_crash_handler, re_grpc_server, re_log, re static GLOBAL: re_memory::AccountingAllocator = re_memory::AccountingAllocator::new(mimalloc::MiMalloc); -/// LCM channel for click events (follows RViz convention) const LCM_CHANNEL: &str = "/clicked_point#geometry_msgs.PointStamped"; -/// Minimum time between click events (debouncing) const CLICK_DEBOUNCE_MS: u64 = 100; -/// Maximum rapid clicks to log as warning -const RAPID_CLICK_THRESHOLD: usize = 5; -/// Default gRPC listen port (9877 to avoid conflict with stock Rerun on 9876) const DEFAULT_PORT: u16 = 9877; -/// DimOS Interactive Viewer — a custom Rerun viewer with LCM click-to-navigate. -/// -/// Accepts the same CLI flags as the stock `rerun` binary so it can be spawned -/// seamlessly via `rerun_bindings.spawn(executable_name="dimos-viewer")`. +/// Entity path prefixes that are considered "robot" entities. +/// Clicking these activates teleop for that robot. +const ROBOT_PREFIXES: &[&str] = &["/world/go2", "/world/g1", "/world/robot"]; + #[derive(Parser, Debug)] #[command(name = "dimos-viewer", version, about)] struct Args { - /// The gRPC port to listen on for incoming SDK connections. #[arg(long, default_value_t = DEFAULT_PORT)] port: u16, - - /// An upper limit on how much memory the viewer should use. - /// When this limit is reached, the oldest data will be dropped. - /// Examples: "75%", "16GB". #[arg(long, default_value = "75%")] memory_limit: String, - - /// An upper limit on how much memory the gRPC server should use. - /// Examples: "1GiB", "50%". #[arg(long, default_value = "1GiB")] server_memory_limit: String, - - /// Hide the Rerun welcome screen. #[arg(long)] hide_welcome_screen: bool, - - /// Hint that data will arrive shortly (suppresses "waiting for data" message). #[arg(long)] expect_data_soon: bool, } -/// Wraps re_viewer::App to add keyboard control interception. +/// Wraps re_viewer::App with keyboard teleop and click-to-nav. struct DimosApp { inner: re_viewer::App, - keyboard: KeyboardHandler, -} - -impl DimosApp { - fn new( - inner: re_viewer::App, - keyboard: KeyboardHandler, - ) -> Self { - Self { - inner, - keyboard, - } - } + keyboard: Rc>, + ctrl_held: Rc>, } impl eframe::App for DimosApp { fn ui(&mut self, ui: &mut egui::Ui, frame: &mut eframe::Frame) { - // Process keyboard input before delegating to Rerun - self.keyboard.process(ui.ctx()); - - // Always draw the keyboard HUD overlay (dims when inactive) - self.keyboard.draw_overlay(ui.ctx()); - - // Delegate to Rerun's main ui method + self.ctrl_held.set(ui.ctx().input(|i| i.modifiers.ctrl || i.modifiers.mac_cmd)); + self.keyboard.borrow_mut().process(ui.ctx()); + self.keyboard.borrow().draw_overlay(ui.ctx()); self.inner.ui(ui, frame); } - // Delegate all other methods to inner re_viewer::App - fn save(&mut self, storage: &mut dyn eframe::Storage) { - self.inner.save(storage); - } - - fn clear_color(&self, visuals: &egui::Visuals) -> [f32; 4] { - self.inner.clear_color(visuals) - } - - fn persist_egui_memory(&self) -> bool { - self.inner.persist_egui_memory() - } - - fn auto_save_interval(&self) -> std::time::Duration { - self.inner.auto_save_interval() - } - + fn save(&mut self, storage: &mut dyn eframe::Storage) { self.inner.save(storage); } + fn clear_color(&self, visuals: &egui::Visuals) -> [f32; 4] { self.inner.clear_color(visuals) } + fn persist_egui_memory(&self) -> bool { self.inner.persist_egui_memory() } + fn auto_save_interval(&self) -> Duration { self.inner.auto_save_interval() } fn raw_input_hook(&mut self, ctx: &egui::Context, raw_input: &mut egui::RawInput) { self.inner.raw_input_hook(ctx, raw_input); } } +/// Check if an entity path belongs to a known robot. +/// Returns the robot prefix (e.g. "/world/go2") if matched. +fn robot_prefix_for(entity_path: &str) -> Option<&'static str> { + ROBOT_PREFIXES.iter().find(|&&p| entity_path.starts_with(p)).copied() +} + #[tokio::main] async fn main() -> Result<(), Box> { let args = Args::parse(); - let main_thread_token = re_viewer::MainThreadToken::i_promise_i_am_on_the_main_thread(); re_log::setup_logging(); re_crash_handler::install_crash_handlers(re_viewer::build_info()); - // Listen for gRPC connections from Rerun's logging SDKs. let listen_addr = format!("0.0.0.0:{}", args.port); re_log::info!("Listening for SDK connections on {listen_addr}"); - let server_memory_limit = re_memory::MemoryLimit::parse(&args.server_memory_limit) - .expect("Bad --server-memory-limit"); let rx_log = re_grpc_server::spawn_with_recv( listen_addr.parse()?, re_grpc_server::ServerOptions { - memory_limit: server_memory_limit, + memory_limit: re_memory::MemoryLimit::parse(&args.server_memory_limit) + .expect("Bad --server-memory-limit"), ..Default::default() }, re_grpc_server::shutdown::never(), ); - // Create LCM publisher for click events let lcm_publisher = LcmPublisher::new(LCM_CHANNEL.to_string()) .expect("Failed to create LCM publisher"); - re_log::info!("LCM publisher created for channel: {LCM_CHANNEL}"); - // Create keyboard handler - let keyboard_handler = KeyboardHandler::new() - .expect("Failed to create keyboard handler"); - re_log::info!("Keyboard handler initialized for WASD controls on /cmd_vel"); + // Shared keyboard handler: DimosApp uses it for process/draw, + // on_event callback uses it to engage/disengage teleop on robot click + let keyboard = Rc::new(RefCell::new( + KeyboardHandler::new().expect("Failed to create keyboard handler") + )); + let keyboard_for_callback = keyboard.clone(); - // State for debouncing and rapid click detection + let ctrl_held = Rc::new(Cell::new(false)); + let ctrl_for_callback = ctrl_held.clone(); let last_click_time = Rc::new(RefCell::new(Instant::now())); - let rapid_click_count = Rc::new(RefCell::new(0usize)); - - let mut native_options = re_viewer::native::eframe_options(None); - native_options.viewport = native_options - .viewport - .with_app_id("rerun_example_custom_callback"); - - let app_env = re_viewer::AppEnvironment::Custom("DimOS Interactive Viewer".to_owned()); let memory_limit = re_memory::MemoryLimit::parse(&args.memory_limit) .expect("Bad --memory-limit"); re_log::info!("Memory limit: {memory_limit}"); + let mut native_options = re_viewer::native::eframe_options(None); + native_options.viewport = native_options.viewport + .with_app_id("rerun_example_custom_callback"); + let startup_options = re_viewer::StartupOptions { memory_limit, - on_event: Some(Rc::new({ - let last_click_time = last_click_time.clone(); - let rapid_click_count = rapid_click_count.clone(); - - move |event: re_viewer::ViewerEvent| { - if let re_viewer::ViewerEventKind::SelectionChange { items } = event.kind { - let mut has_position = false; - let mut no_position_count = 0; + on_event: Some(Rc::new(move |event: re_viewer::ViewerEvent| { + if let re_viewer::ViewerEventKind::SelectionChange { items } = event.kind { + for item in &items { + if let re_viewer::SelectionChangeItem::Entity { entity_path, position, .. } = item { + let path = entity_path.to_string(); + + // Check if clicked entity is a robot → engage teleop + if let Some(prefix) = robot_prefix_for(&path) { + let mut kb = keyboard_for_callback.borrow_mut(); + kb.set_active_robot(Some(prefix.to_string())); + kb.set_engaged(true); + re_log::info!("Teleop engaged: {prefix}"); + return; // Robot click = engage only, not nav goal + } - for item in items { - match item { - re_viewer::SelectionChangeItem::Entity { - entity_path, - view_name: _, - position: Some(pos), - .. - } => { - has_position = true; + // Not a robot entity: disengage teleop + { + let mut kb = keyboard_for_callback.borrow_mut(); + if kb.engaged() { + kb.set_engaged(false); + re_log::info!("Teleop disengaged"); + } + } - // Debouncing + // Ctrl+click on non-robot entity with position → nav goal + if ctrl_for_callback.get() { + if let Some(pos) = position { let now = Instant::now(); - let elapsed = now.duration_since(*last_click_time.borrow()); - - if elapsed < Duration::from_millis(CLICK_DEBOUNCE_MS) { - let mut count = rapid_click_count.borrow_mut(); - *count += 1; - if *count == RAPID_CLICK_THRESHOLD { - re_log::warn!( - "Rapid click detected ({} clicks within {}ms)", - RAPID_CLICK_THRESHOLD, - CLICK_DEBOUNCE_MS - ); - } + if now.duration_since(*last_click_time.borrow()) < Duration::from_millis(CLICK_DEBOUNCE_MS) { continue; - } else { - *rapid_click_count.borrow_mut() = 0; } *last_click_time.borrow_mut() = now; - let timestamp_ms = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64; - - // Build click event and publish via LCM - let click = click_event_from_ms( - [pos.x, pos.y, pos.z], - &entity_path.to_string(), - timestamp_ms, - ); - - match lcm_publisher.publish(&click) { - Ok(_) => { - re_log::debug!( - "LCM click event published: entity={}, pos=({:.2}, {:.2}, {:.2})", - entity_path, - pos.x, - pos.y, - pos.z - ); - } - Err(err) => { - re_log::error!("Failed to publish LCM click event: {err:?}"); - } + let ts = SystemTime::now().duration_since(UNIX_EPOCH) + .unwrap_or_default().as_millis() as u64; + let click = click_event_from_ms([pos.x, pos.y, pos.z], "world", ts); + if let Err(e) = lcm_publisher.publish(&click) { + re_log::error!("Nav goal failed: {e:?}"); + } else { + re_log::info!("Nav goal: ({:.2}, {:.2}, {:.2})", pos.x, pos.y, pos.z); } } - re_viewer::SelectionChangeItem::Entity { position: None, .. } => { - no_position_count += 1; - } - _ => {} } } - - if !has_position && no_position_count > 0 { - re_log::trace!( - "Selection change without position data ({no_position_count} items). \ - This is normal for hover/keyboard navigation." - ); - } } } })), ..Default::default() }; - let window_title = "DimOS Interactive Viewer"; eframe::run_native( - window_title, + "DimOS Interactive Viewer", native_options, Box::new(move |cc| { re_viewer::customize_eframe_and_setup_renderer(cc)?; - let mut rerun_app = re_viewer::App::new( main_thread_token, re_viewer::build_info(), - app_env, + re_viewer::AppEnvironment::Custom("DimOS Interactive Viewer".to_owned()), startup_options, cc, None, re_viewer::AsyncRuntimeHandle::from_current_tokio_runtime_or_wasmbindgen()?, ); - rerun_app.add_log_receiver(rx_log); - - let dimos_app = DimosApp::new(rerun_app, keyboard_handler); - - Ok(Box::new(dimos_app)) + Ok(Box::new(DimosApp { inner: rerun_app, keyboard, ctrl_held })) }), )?; - Ok(()) }