assigned control joints. needs testing

OSUrobotics · Nov 14, 2024 · a948891 · a948891
1 parent 348ba2a
commit a948891
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/pybullet_tree_sim/robot.py b/pybullet_tree_sim/robot.py
@@ -64,6 +64,7 @@ def __init__(
 
         self.robot_conf = {}
         self._generate_robot_urdf()
+        self._assign_control_joints()
         self.setup_robot()
 
         return
@@ -93,9 +94,25 @@ def _generate_robot_urdf(self) -> None:
             robot_urdf = robot_urdf.toprettyxml()
 
         # Save the generated URDF
+        self.robot_urdf: str = robot_urdf
         self.robot_urdf_path = os.path.join(self._urdf_tmp_path, "robot.urdf")
         xutils.save_urdf(robot_urdf, urdf_path=self.robot_urdf_path)
         return
+
+    def _assign_control_joints(self) -> None:
+        self.robot_conf['control_joints'] = []
+        with open(self.robot_urdf_path) as f:
+            for line in f:
+                line = line.strip()
+                if line.startswith('<joint') and "type=" in line:
+                    log.warn(line)
+                    joint_name = line.split('name="')[1].split('"')[0]
+                    joint_type = line.split('type="')[1].split('"')[0]
+                    if joint_type == 'fixed':
+                        continue
+                    else:
+                        self.robot_conf['control_joints'].append(joint_name)
+        return
 
     def setup_robot(self):
         if self.robot is not None:
@@ -155,7 +172,7 @@ def setup_robot(self):
                 if self.verbose > 1:
                     print("Joint Name: ", jointName, "Joint ID: ", jointID)
 
-                controllable = True if jointName in self.control_joints else False
+                controllable = True if jointName in self.robot_conf["control_joints"] else False
                 if controllable:
                     self.controllable_joints_idxs.append(i)
                     self.joint_lower_limits.append(jointLowerLimit)