Added ability to export the frames of the robot during solve

README.md

@@ -4,6 +4,8 @@
 
 ## About
 
+LivelyTK Package
+
 The LivelyTK framework provides a highly configurable toolkit for commanding robots in mixed modalities while incorporating liveliness motions. It is adapted from [RelaxedIK](https://github.com/uwgraphics/relaxed_ik_core) framework.
 
 To configure a robot, the easiest method is to use the LivelyStudio interface in the [lively_tk_ros](https://github.com/Wisc-HCI/lively_tk_ros) repository, which is a wizard for configuring the robot.
@@ -36,17 +38,18 @@ inputs = [
     ObjectiveInput(weight=3.0),
     ObjectiveInput(weight=1.0,quaternion=[1,0,0,0])
 ]
-base_position, joint_values = solver.solve(inputs,timestamp,max_retries=2)
+base_position, joint_values, _ = solver.solve(inputs,timestamp,max_retries=2)
 
 ```
 
 The fields for solve are as follows:
 1. `goals`: a list of ObjectiveInput objects.
-2. `time`: a float indicating the current time. If no liveliness objectives are used, this has no effect.
+2. `time`: (float) The current time. If no liveliness objectives are used, this has no effect.
 3. `world`: [NOT IMPLEMENTED]. Update the current model of the world to handle real-time collision avoidance.
-4. `max_retries`: Number of random restarts to perform when searching for a solution.
-5. `max_iterations`: Number of iterations for each round of search.
-6. `only_core`: Ignore liveliness objectives and disable the second liveliness optimization. *Note: Not advised to switch within runs. Only use this flag if you are not using liveliness objectives and want a slight speed-up.*
+4. `max_retries`: (int) Number of random restarts to perform when searching for a solution.
+5. `max_iterations`: (int) Number of iterations for each round of search.
+6. `only_core`: (bool) Ignore liveliness objectives and disable the second liveliness optimization. *Note: Not advised to switch within runs. Only use this flag if you are not using liveliness objectives and want a slight speed-up.*
+7. `return_frames`: (bool) Return the positions and orientations of each joint (structured specified in joint_names in config) as the third output field.
 
 If you need to restart a run, you can execute the reset method of the solver, and provide the starting base position and joint values:
 

setup.py

@@ -38,17 +38,18 @@
     ObjectiveInput(weight=3.0),
     ObjectiveInput(weight=1.0,quaternion=[1,0,0,0])
 ]
-base_position, joint_values = solver.solve(inputs,timestamp,max_retries=2)
+base_position, joint_values, _ = solver.solve(inputs,timestamp,max_retries=2)
 
 ```
 
 The fields for solve are as follows:
 1. `goals`: a list of ObjectiveInput objects.
-2. `time`: a float indicating the current time. If no liveliness objectives are used, this has no effect.
+2. `time`: (float) The current time. If no liveliness objectives are used, this has no effect.
 3. `world`: [NOT IMPLEMENTED]. Update the current model of the world to handle real-time collision avoidance.
-4. `max_retries`: Number of random restarts to perform when searching for a solution.
-5. `max_iterations`: Number of iterations for each round of search.
-6. `only_core`: Ignore liveliness objectives and disable the second liveliness optimization. *Note: Not advised to switch within runs. Only use this flag if you are not using liveliness objectives and want a slight speed-up.*
+4. `max_retries`: (int) Number of random restarts to perform when searching for a solution.
+5. `max_iterations`: (int) Number of iterations for each round of search.
+6. `only_core`: (bool) Ignore liveliness objectives and disable the second liveliness optimization. *Note: Not advised to switch within runs. Only use this flag if you are not using liveliness objectives and want a slight speed-up.*
+7. `return_frames`: (bool) Return the positions and orientations of each joint (structured specified in joint_names in config) as the third output field.
 
 If you need to restart a run, you can execute the reset method of the solver, and provide the starting base position and joint values:
 
@@ -59,7 +60,7 @@
 
 setup(
     name=package_name,
-    version='0.7.0',
+    version='0.8.0',
     packages=[package_name],
     rust_extensions=[RustExtension("lively_tk.lively_tk", binding=Binding.PyO3, quiet=True)],
     install_requires=['setuptools','wheel','setuptools_rust'],

src/lively_tk.rs

@@ -74,8 +74,9 @@ impl Solver {
         world: Option<EnvironmentSpec>,
         max_retries: Option<u64>,
         max_iterations: Option<usize>,
-        only_core: Option<bool>
-    ) -> PyResult<(Vec<f64>, Vec<f64>)> {
+        only_core: Option<bool>,
+        return_frames: Option<bool>
+    ) -> PyResult<(Vec<f64>, Vec<f64>, Option<Vec<Vec<(Vec<f64>, Vec<f64>)>>>)> {
         // RNG Gen
         let mut rng = thread_rng();
 
@@ -93,6 +94,11 @@ impl Solver {
             Some(v) => run_only_core = v,
             None => {}
         }
+        let mut provide_frames = false;
+        match return_frames {
+            Some(v) => provide_frames = v,
+            None => {}
+         }
 
         for i in 0..out_x.len() {
             xopt.push(out_x[i])
@@ -108,7 +114,7 @@ impl Solver {
                 goals.len(),
                 self.config.objectives.len()
             );
-            return Ok((self.vars.offset.clone(), out_x));
+            return Ok((self.vars.offset.clone(), out_x, None));
         }
 
         for goal_idx in 0..goals.clone().len() {
@@ -233,7 +239,33 @@ impl Solver {
         }
 
         // println!("OUTX-CORE {:?},\nOUTX {:?}",xopt_core,xopt);
+        if provide_frames {
+            let raw_frames = self.vars.robot.get_frames(&self.vars.history.prev1.clone());
+            let mut return_frame_values: Vec<Vec<(Vec<f64>, Vec<f64>)>> = Vec::new();
+            for arm_idx in 0..raw_frames.len() {
+                let mut arm_values: Vec<(Vec<f64>, Vec<f64>)> = Vec::new();
+                for joint_idx in 0..raw_frames[arm_idx].0.len() {
+                    arm_values.push((
+                        vec![
+                            raw_frames[arm_idx].0[joint_idx].x,
+                            raw_frames[arm_idx].0[joint_idx].y,
+                            raw_frames[arm_idx].0[joint_idx].z
+                        ],
+                        vec![
+                            raw_frames[arm_idx].1[joint_idx].as_vector()[3],
+                            raw_frames[arm_idx].1[joint_idx].as_vector()[0],
+                            raw_frames[arm_idx].1[joint_idx].as_vector()[1],
+                            raw_frames[arm_idx].1[joint_idx].as_vector()[2]
+                        ]
+                    ))
+                }
+                return_frame_values.push(arm_values);
+            }
+            return Ok((self.vars.offset.clone(), self.vars.xopt.clone(), Some(return_frame_values)));
+        } else {
+            return Ok((self.vars.offset.clone(), self.vars.xopt.clone(), None));
+        }
 
-        return Ok((self.vars.offset.clone(), self.vars.xopt.clone()));
+
     }
 }