Solidify Scene{Description,State} typing (#2147)

* Expose Rust scene types to Python

* sort imports

* isort

* isort

* remove unused imports

Author: oleflb

Cargo.lock

@@ -18,10 +18,10 @@ use hardware::{
     TransformMessageInterface,
 };
 use hula_types::hardware::{Ids, Paths};
-use log::{error, warn};
+use log::{error, info, warn};
 use parking_lot::Mutex;
 use serde::Deserialize;
-use simulation_message::{ClientMessageKind, ConnectionInfo, ServerMessageKind, SimulationMessage};
+use simulation_message::{ClientMessageKind, ConnectionInfo, ServerMessageKind, SimulatorMessage};
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::time::sleep;
 use tokio_tungstenite::tungstenite::Message;
@@ -123,9 +123,7 @@ async fn worker(
             log::info!("connected to mujoco websocket at {address}");
             log::info!("sending ConnectionInfo");
             websocket
-                .send(Message::Text(
-                    serde_json::to_string(&connection_info)?.into(),
-                ))
+                .send(Message::binary(bincode::serialize(&connection_info)?))
                 .await?;
             break websocket;
         };
@@ -144,7 +142,7 @@ async fn worker(
             },
             maybe_low_command_event = worker_channels.low_command_receiver.recv() => {
                 match maybe_low_command_event {
-                    Some(low_command) => websocket.send(Message::Text(serde_json::to_string(&ClientMessageKind::LowCommand(low_command))?.into())).await?,
+                    Some(low_command) => websocket.send(Message::binary(bincode::serialize(&ClientMessageKind::LowCommand(low_command))?)).await?,
                     None => break,
                 };
             },
@@ -161,22 +159,22 @@ async fn handle_message(
     worker_channels: &WorkerChannels,
 ) -> Result<()> {
     let message = match message {
-        Message::Text(string) => serde_json::from_str(&string)?,
+        Message::Binary(data) => bincode::deserialize(&data)?,
         Message::Close(maybe_frame) => {
             warn!("server closed connections: {maybe_frame:#?}");
             return Ok(());
         }
         _ => return Ok(()),
     };
     match message {
-        SimulationMessage {
+        SimulatorMessage {
             payload: ServerMessageKind::LowState(low_state),
             time,
         } => {
             *hardware_interface_time.lock() = time;
             worker_channels.low_state_sender.send(low_state).await?
         }
-        SimulationMessage {
+        SimulatorMessage {
             payload: ServerMessageKind::FallDownState(fall_down_state),
             time,
         } => {
@@ -186,7 +184,7 @@ async fn handle_message(
                 .send(fall_down_state)
                 .await?
         }
-        SimulationMessage {
+        SimulatorMessage {
             payload: ServerMessageKind::ButtonEventMsg(button_event_msg),
             time,
         } => {
@@ -196,7 +194,7 @@ async fn handle_message(
                 .send(button_event_msg)
                 .await?
         }
-        SimulationMessage {
+        SimulatorMessage {
             payload: ServerMessageKind::RemoteControllerState(remote_controller_state),
             time,
         } => {
@@ -206,7 +204,7 @@ async fn handle_message(
                 .send(remote_controller_state)
                 .await?
         }
-        SimulationMessage {
+        SimulatorMessage {
             payload: ServerMessageKind::TransformMessage(transform_stamped),
             time,
         } => {
@@ -216,7 +214,7 @@ async fn handle_message(
                 .send(transform_stamped)
                 .await?
         }
-        SimulationMessage {
+        SimulatorMessage {
             payload: ServerMessageKind::RGBDSensors(rgbd_sensors),
             time,
         } => {
@@ -226,6 +224,9 @@ async fn handle_message(
                 .send(*rgbd_sensors)
                 .await?
         }
+        _ => {
+            info!("Received unexpected simulator data")
+        }
     };
 
     Ok(())

crates/hulk_mujoco/src/hardware_interface.rs

@@ -1,17 +1,18 @@
 use std::{
+    collections::BTreeMap,
     ops::Range,
     time::{Duration, SystemTime},
 };
 
 use booster::{
     ButtonEventMsg, FallDownState, LowCommand, LowState, RemoteControllerState, TransformMessage,
 };
-use pyo3::pyclass;
+use pyo3::{pyclass, pymethods};
 use serde::{Deserialize, Serialize};
 use zed::RGBDSensors;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct SimulationMessage<T> {
+pub struct SimulatorMessage<T> {
     pub time: SystemTime,
     pub payload: T,
 }
@@ -24,6 +25,8 @@ pub enum ServerMessageKind {
     RemoteControllerState(RemoteControllerState),
     TransformMessage(TransformMessage),
     RGBDSensors(Box<RGBDSensors>),
+    SceneUpdate(SceneUpdate),
+    SceneDescription(SceneDescription),
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -44,6 +47,137 @@ pub enum TaskName {
     RequestSceneDescription,
 }
 
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct SceneDescription {
+    pub meshes: BTreeMap<String, SceneMesh>,
+    pub lights: Vec<Light>,
+    pub bodies: BTreeMap<String, Body>,
+}
+
+#[pymethods]
+impl SceneDescription {
+    #[new]
+    pub fn new(
+        meshes: BTreeMap<String, SceneMesh>,
+        lights: Vec<Light>,
+        bodies: BTreeMap<String, Body>,
+    ) -> Self {
+        Self {
+            meshes,
+            lights,
+            bodies,
+        }
+    }
+}
+
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct SceneMesh {
+    pub vertices: Vec<[f32; 3]>,
+    pub faces: Vec<[u32; 3]>,
+}
+
+#[pymethods]
+impl SceneMesh {
+    #[new]
+    pub fn new(vertices: Vec<[f32; 3]>, faces: Vec<[u32; 3]>) -> Self {
+        Self { vertices, faces }
+    }
+}
+
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Light {
+    pub name: Option<String>,
+    pub pos: [f32; 3],
+    pub dir: [f32; 3],
+}
+
+#[pymethods]
+impl Light {
+    #[new]
+    pub fn new(name: Option<String>, pos: [f32; 3], dir: [f32; 3]) -> Self {
+        Self { name, pos, dir }
+    }
+}
+
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Body {
+    pub id: i64,
+    pub parent: Option<String>,
+    pub geoms: Vec<Geom>,
+}
+
+#[pymethods]
+impl Body {
+    #[new]
+    pub fn new(id: i64, parent: Option<String>, geoms: Vec<Geom>) -> Self {
+        Self { id, parent, geoms }
+    }
+}
+
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Geom {
+    pub name: Option<String>,
+    pub mesh: Option<String>,
+    pub rgba: [f32; 4],
+    pub pos: [f32; 3],
+    pub quat: [f32; 4],
+}
+
+#[pymethods]
+impl Geom {
+    #[new]
+    pub fn new(
+        name: Option<String>,
+        mesh: Option<String>,
+        rgba: [f32; 4],
+        pos: [f32; 3],
+        quat: [f32; 4],
+    ) -> Self {
+        Self {
+            name,
+            mesh,
+            rgba,
+            pos,
+            quat,
+        }
+    }
+}
+
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct SceneUpdate {
+    pub time: f32,
+    pub bodies: BTreeMap<String, BodyUpdate>,
+}
+
+#[pymethods]
+impl SceneUpdate {
+    #[new]
+    pub fn new(time: f32, bodies: BTreeMap<String, BodyUpdate>) -> Self {
+        Self { time, bodies }
+    }
+}
+
+#[pyclass(frozen)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct BodyUpdate {
+    pub pos: [f32; 3],
+    pub quat: [f32; 4],
+}
+
+#[pymethods]
+impl BodyUpdate {
+    #[new]
+    pub fn new(pos: [f32; 3], quat: [f32; 4]) -> Self {
+        Self { pos, quat }
+    }
+}
+
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ConnectionInfo {
     pub schedule: Vec<TaskSchedule>,

crates/simulation_message/src/lib.rs

@@ -10,6 +10,7 @@ name = "mujoco_rust_server"
 crate-type = ["cdylib"]
 
 [dependencies]
+bincode = { workspace = true }
 booster = { workspace = true }
 color-eyre = { workspace = true }
 futures-util = { workspace = true }

tools/mujoco-simulator/mujoco-rust-server/Cargo.toml

@@ -12,7 +12,11 @@ class PySimulationTask:
     async def respond(
         self,
         time: float,
-        response: booster_types.LowState | zed_types.RGBDSensors | bytes | str | None,
+        response: booster_types.LowState
+        | zed_types.RGBDSensors
+        | bytes
+        | str
+        | None,
     ) -> None: ...
     async def receive(self) -> booster_types.LowCommand: ...
 
@@ -26,8 +30,49 @@ class TaskName(Enum):
     Reset = auto()
     Invalid = auto()
 
+class Body:
+    id: int
+    parent: str | None
+    geoms: list[Geom]
+
+class BodyUpdate:
+    pos: list[float]
+    quat: list[float]
+
+class Geom:
+    name: str
+    mesh: str | None
+    rgba: list[float]
+    pos: list[float]
+    quat: list[float]
+
+class Light:
+    name: str | None
+    pos: list[float]
+    dir: list[float]
+
+class SceneDescription:
+    meshes: dict[str, SceneMesh]
+    lights: list[Light]
+    bodies: dict[str, Body]
+
+class SceneMesh:
+    vertices: list[list[float]]
+    faces: list[list[int]]
+
+class SceneUpdate:
+    time: float
+    bodies: dict[str, BodyUpdate]
+
 __all__ = [
+    "Body",
+    "BodyUpdate",
+    "Geom",
+    "Light",
     "PySimulationTask",
+    "SceneDescription",
+    "SceneMesh",
+    "SceneUpdate",
     "SimulationServer",
     "TaskName",
     "booster_types",