expose more methods (#47)

kscalelabs · Nov 11, 2024 · 40b4683 · 40b4683
1 parent 508a9ac
commit 40b4683
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 20 deletions.
diff --git a/actuator/bindings/src/lib.rs b/actuator/bindings/src/lib.rs
@@ -258,6 +258,27 @@ impl PyRobstrideMotorsSupervisor {
         Ok(PyRobstrideMotorsSupervisor { inner: controller })
     }
 
+    fn set_all_params(&self, params: HashMap<u8, PyRobstrideMotorControlParams>) -> PyResult<()> {
+        let params: HashMap<u8, RobstrideMotorControlParams> = params
+            .into_iter()
+            .map(|(k, v)| (k, RobstrideMotorControlParams::from(v)))
+            .collect();
+        self.inner.set_all_params(params);
+        Ok(())
+    }
+
+    fn set_params(&self, motor_id: u8, params: PyRobstrideMotorControlParams) -> PyResult<()> {
+        self.inner
+            .set_params(motor_id, RobstrideMotorControlParams::from(params))
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+    }
+
+    fn set_positions(&self, positions: HashMap<u8, f32>) -> PyResult<()> {
+        self.inner
+            .set_positions(positions)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+    }
+
     fn set_position(&self, motor_id: u8, position: f32) -> PyResult<f32> {
         self.inner
             .set_position(motor_id, position)
@@ -270,6 +291,12 @@ impl PyRobstrideMotorsSupervisor {
             .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
     }
 
+    fn set_velocities(&self, velocities: HashMap<u8, f32>) -> PyResult<()> {
+        self.inner
+            .set_velocities(velocities)
+            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
+    }
+
     fn set_velocity(&self, motor_id: u8, velocity: f32) -> PyResult<f32> {
         self.inner
             .set_velocity(motor_id, velocity)
@@ -350,20 +377,6 @@ impl PyRobstrideMotorsSupervisor {
         ))
     }
 
-    fn set_params(&self, motor_id: u8, params: &PyRobstrideMotorControlParams) -> PyResult<()> {
-        self.inner.set_params(
-            motor_id,
-            RobstrideMotorControlParams {
-                position: params.position,
-                velocity: params.velocity,
-                kp: params.kp,
-                kd: params.kd,
-                torque: params.torque,
-            },
-        );
-        Ok(())
-    }
-
     #[getter]
     fn total_commands(&self) -> PyResult<u64> {
         Ok(self.inner.get_total_commands())

diff --git a/actuator/robstride/src/supervisor.rs b/actuator/robstride/src/supervisor.rs
@@ -156,14 +156,17 @@ impl MotorsSupervisor {
                 }
 
                 {
-                    // Send PD commands to motors.
-                    let target_params = target_params.read().unwrap();
-                    if !target_params.is_empty() {
-                        match motors.send_motor_controls(&target_params, *serial.read().unwrap()) {
+                    let params_copy = {
+                        let target_params = target_params.read().unwrap();
+                        target_params.clone()
+                    };
+
+                    if !params_copy.is_empty() {
+                        match motors.send_motor_controls(&params_copy, *serial.read().unwrap()) {
                             Ok(feedbacks) => {
                                 let mut latest_feedback = latest_feedback.write().unwrap();
                                 let mut failed_commands = failed_commands.write().unwrap();
-                                for &motor_id in target_params.keys() {
+                                for &motor_id in params_copy.keys() {
                                     if let Some(feedback) = feedbacks.get(&motor_id) {
                                         latest_feedback.insert(motor_id, feedback.clone());
                                     } else {
@@ -236,9 +239,27 @@ impl MotorsSupervisor {
         *self.failed_commands.write().unwrap() = HashMap::new();
     }
 
-    pub fn set_params(&self, motor_id: u8, params: MotorControlParams) {
+    pub fn set_all_params(&self, params: HashMap<u8, MotorControlParams>) {
+        let mut target_params = self.target_params.write().unwrap();
+        *target_params = params;
+    }
+
+    pub fn set_params(
+        &self,
+        motor_id: u8,
+        params: MotorControlParams,
+    ) -> Result<(), std::io::Error> {
         let mut target_params = self.target_params.write().unwrap();
         target_params.insert(motor_id, params);
+        Ok(())
+    }
+
+    pub fn set_positions(&self, positions: HashMap<u8, f32>) -> Result<(), std::io::Error> {
+        let mut target_params = self.target_params.write().unwrap();
+        for (motor_id, position) in positions {
+            target_params.get_mut(&motor_id).unwrap().position = position;
+        }
+        Ok(())
     }
 
     pub fn set_position(&self, motor_id: u8, position: f32) -> Result<f32, std::io::Error> {
@@ -267,6 +288,14 @@ impl MotorsSupervisor {
             })
     }
 
+    pub fn set_velocities(&self, velocities: HashMap<u8, f32>) -> Result<(), std::io::Error> {
+        let mut target_params = self.target_params.write().unwrap();
+        for (motor_id, velocity) in velocities {
+            target_params.get_mut(&motor_id).unwrap().velocity = velocity;
+        }
+        Ok(())
+    }
+
     pub fn set_velocity(&self, motor_id: u8, velocity: f32) -> Result<f32, std::io::Error> {
         let mut target_params = self.target_params.write().unwrap();
         if let Some(params) = target_params.get_mut(&motor_id) {