updated main, penv, robot to work with tree refactor

Author: lukestroh

main.py

@@ -15,7 +15,7 @@
 
 def main():
     pbutils = PyBUtils(renders=True)
-    robot_start_orientation = Rotation.from_euler('xyz', [0, 0,180], degrees=True).as_quat()
+    robot_start_orientation = Rotation.from_euler("xyz", [0, 0, 180], degrees=True).as_quat()
     robot = Robot(pbclient=pbutils.pbclient, position=[0, 1, 0], orientation=robot_start_orientation)
 
     penv = PruningEnv(
@@ -42,7 +42,7 @@ def main():
         save_tree_urdf=False,
         # randomize_pose=True
     )
-    penv.activate_tree(tree_id_str=tree_name)
+    penv.activate_tree_by_id_str(tree_id_str=tree_name)
 
     # # Run the sim a little just to get the environment properly loaded.
     for i in range(100):

pybullet_tree_sim/pruning_environment.py

@@ -69,7 +69,7 @@ def __init__(
         # make_trees: bool = False,
         name: str = "PruningEnv",
         # num_trees: int | None = None,
-        renders: bool = False,
+        # renders: bool = False,
         verbose: bool = True,
     ) -> None:
         """Initialize the Pruning Environment
@@ -83,7 +83,7 @@ def __init__(
 
         # Pybullet GUI variables
         self.render_mode = "rgb_array"
-        self.renders = renders
+        # self.renders = renders
         # self.eval = evaluate
 
         # Gym variables
@@ -143,61 +143,54 @@ def load_tree(  # TODO: Clean up Tree init vs create_tree, probably not needed.
             randomize_pose=randomize_pose,
         )
 
-        tree_id_str = f"{tree_namespace}{tree_type}_tree{tree_id}"
-        # urdf_path = os.path.join(URDF_PATH, "trees", tree_type, "generated", f"{tree.id_str}.urdf")
-
         # Add tree to dict of trees
         self.trees[tree.id_str] = tree
-        return tree_id_str
+        return tree.id_str
+
+    def get_tree_from_id_str(self, tree_id_str: str) -> Tree:
+        try:
+            tree = self.trees[tree_id_str]
+        except KeyError as e:
+            raise TreeException(f"{e}: Tree with ID {tree_id_str} not found")
+        return tree
 
     def activate_tree(
         self,
-        tree: Tree | None = None,
-        tree_id_str: str | None = None,
+        tree: Tree,
         include_support_posts: bool = True,
     ) -> None:
         """Activate a tree by object or by tree_id_str. Can include support posts. Must provide either a Tree or tree_id_str.
         @param tree (Tree/None): Tree object to be activated into the pruning environment.
         @param tree_id_str (str/None): String including the identification characteristics of the tree.
         @return None
         """
-
-        if tree is None and tree_id_str is None:
-            raise TreeException("Parameters 'tree' and 'tree_id_str' cannot both be None")
-
-        if tree is None and tree_id_str is not None:
-            try:
-                tree = self.trees[tree_id_str]
-            except KeyError as e:
-                raise TreeException(f"{e}: Tree with ID {tree_id_str} not found")
-
-        if tree is not None:
+        if tree:
             if self.verbose:
                 log.info("Activating tree")
-            tree.pyb_tree_id = self.pbutils.pbclient.loadURDF(tree.urdf_path, useFixedBase=True)
-            log.info(f"Tree {tree.id_str} activated with PyBID {tree.pyb_tree_id}")
+            tree.pyb_id = self.pbutils.pbclient.loadURDF(tree.urdf_path, useFixedBase=True)
+            log.info(f"Tree {tree.id_str} activated with PyBID {tree.pyb_id}")
 
             if include_support_posts:
                 self.activate_support_posts(associated_tree=tree)
         return
 
-    def deactivate_tree(self, tree: Tree | None = None, tree_id_str: str | None = None) -> None:
-        """Deactivate a tree by object or by tree_id_str"""
-        if tree is None and tree_id_str is None:
-            raise TreeException("Parameters 'tree' and 'tree_id_str' cannot both be None")
-
-        if tree is None and tree_id_str is not None:
-            try:
-                tree = self.trees[tree_id_str]
-            except KeyError as e:
-                raise TreeException(f"{e}: Tree with ID {tree_id_str} not found")
-
-        if tree is not None:
-            try:
-                self.pbutils.pbclient.removeBody(tree.pyb_tree_id)
-                log.info(f"Tree {tree.id_str} with PyBID {tree.pyb_tree_id} deactivated")
-            except Exception as e:
-                log.error(f"Error deactivating tree: {e}")
+    def activate_tree_by_id_str(self, tree_id_str: str, include_support_posts: bool = True) -> None:
+        tree = self.get_tree_from_id_str(tree_id_str=tree_id_str)
+        self.activate_tree(tree=tree, include_support_posts=include_support_posts)
+        return
+
+    def deactivate_tree(self, tree: Tree) -> None:
+        """Deactivate a tree object"""
+        try:
+            self.pbutils.pbclient.removeBody(tree.pyb_id)
+            log.info(f"Tree {tree.id_str} with PyBID {tree.pyb_id} deactivated")
+        except Exception as e:
+            log.error(f"Error deactivating tree: {e}")
+        return
+
+    def deactivate_tree_by_id_str(self, tree_id_str: str) -> None:
+        tree = self.get_tree_from_id_str(tree_id_str=tree_id_str)
+        self.deactivate_tree(tree=tree)
         return
 
     def activate_support_posts(

pybullet_tree_sim/robot.py

@@ -44,7 +44,7 @@ def __init__(
         self.position = position
         self.orientation = orientation
         self.randomize_pose = randomize_pose  # TODO: This isn't set up anymore... fix
-        self.init_joint_angles = (
+        self.init_joint_angles = (  # TODO: dynamically assign from defaults (i.e. if linear-slider is loaded in)
             (
                 -np.pi / 2 + np.pi / 4,
                 -np.pi * 2 / 3,
@@ -172,6 +172,7 @@ def _setup_robot(self):
     def _get_joints(self) -> dict:
         """Return a dict of joint information for the robot"""
         joints = {}
+
         for i in range(self.num_joints):
             info = self.pbclient.getJointInfo(self.robot, i)
             joint_name = info[1].decode("utf-8")
@@ -187,16 +188,25 @@ def _get_joints(self) -> dict:
                     }
                 }
             )
+        log.warn(joints)
         return joints
 
-    def _assign_control_joints(self, joints: dict) -> list:
+    def _assign_control_joints(self, joints: dict) -> tuple[list]:
         """Get list of controllable joints from the joint dict by joint type"""
         control_joints = []
         control_joint_idxs = []
+        self.control_joint_lower_limits = []
+        self.control_joint_upper_limits = []
+        self.control_joint_ranges = []
         for joint, joint_info in joints.items():
             if joint_info["type"] == 0:  # TODO: Check if this works for prismatic joints or just revolute
                 control_joints.append(joint)
                 control_joint_idxs.append(joint_info["id"])
+
+                # self.joint_upper_limits,
+                # self.joint_lower_limits,
+                # self.joint_ranges,  # ,
+
         return control_joints, control_joint_idxs
 
     def _get_links(self) -> dict:
@@ -323,6 +333,7 @@ def reset_robot(self, joint_angles: tuple = None) -> None:
             else joint_angles
         )
         self.set_joint_angles_no_collision(self.init_joint_angles)
+        # self.pbclient.resetJointState() # TODO Fill out
         return
 
     def remove_robot(self):
@@ -565,7 +576,7 @@ def check_collisions(self, collision_objects) -> Tuple[bool, dict]:
             collisons_self = self.pbclient.getContactPoints(bodyA=self.robot, bodyB=self.robot)
             collisions_unacceptable = collisons_self
             for i in range(len(collisions_unacceptable)):
-                if collisions_unacceptable[i][-6] < -0.001:
+                if collisions_unacceptable[i][-6] < -0.00:
                     collision_info["collisions_unacceptable"] = True
                     break
         if self.verbose > 1:
@@ -578,7 +589,7 @@ def check_collisions(self, collision_objects) -> Tuple[bool, dict]:
 
     def check_success_collision(self, body_b) -> bool:
         """Check if there are any collisions between the robot and the environment
-        Returns: Boolw
+        Returns: Bool
         """
         collisions_success = self.pbclient.getContactPoints(
             bodyA=self.robot, bodyB=body_b, linkIndexA=self.success_link_index