current updates for cindy code review

Author: lukestroh

main.py

@@ -2,6 +2,7 @@
 from pybullet_tree_sim.camera import Camera
 from pybullet_tree_sim.time_of_flight import TimeOfFlight
 from pybullet_tree_sim.pruning_environment import PruningEnv
+from pybullet_tree_sim.robot import Robot
 from pybullet_tree_sim.tree import Tree
 from pybullet_tree_sim.utils.pyb_utils import PyBUtils
 from pybullet_tree_sim.utils.camera_helpers import get_fov_from_dfov
@@ -20,13 +21,27 @@ def main():
         "camera0": {
             'sensor': Camera(pbutils, sensor_name="realsense_d435i"),
             'tf_frame': "mock_pruner__camera0",
+        },
+        'tof0': {
+            'sensor': TimeOfFlight(pbutils, sensor_name="vl6180"),
+            'tf_frame': "mock_pruner__tof0",
+        },
+        'tof1': {
+            'sensor': TimeOfFlight(pbutils, sensor_name="vl6180"),
+            'tf_frame': "mock_pruner__tof1",
         }
     }
 
+    robots = {
+        'pruning_robot': Robot(
+            pbclient=pbutils.pbclient,
+            sensor_config=sensor_config,
+            position = [0, 1, 0],
+            orientation = [0, 0, 0, 1])
+    }
 
-    
     penv = PruningEnv(
-        pbutils=pbutils, load_robot=True, sensor_config=sensor_config, robot_pos=[0, 2, 0], verbose=True,
+        pbutils=pbutils, robots=robots, verbose=True,
     )
 
 
@@ -48,25 +63,25 @@ def main():
         pbutils.pbclient.stepSimulation()
         time.sleep(0.1)
 
-    time.sleep(30)
-
+    # time.sleep(30)
+    robot = penv.robots['pruning_robot']
     # Simulation loop
     while True:
         try:
-            tof0_view_matrix = penv.robot.get_view_mat_at_curr_pose(camera=penv.robot.sensors['mock_pruner__camera0'])
-            rgbd = penv.pbutils.get_rgbd_at_cur_pose(type="robot", view_matrix=view_matrix)
+            sensor_view_matrix = robot.get_view_mat_at_curr_pose(camera=robot.sensors['mock_pruner__camera0'])
+            rgbd = penv.pbutils.get_rgbd_at_cur_pose(camera=robot.sensors['mock_pruner__camera0'], type="robot", view_matrix=sensor_view_matrix)
             keys_pressed = penv.get_key_pressed()
-            action = penv.get_key_action(keys_pressed=keys_pressed)
+            action = penv.get_key_action(robot=robot, keys_pressed=keys_pressed)
             action = action.reshape((6, 1))
-            jv, jacobian = penv.robot.calculate_joint_velocities_from_ee_velocity_dls(end_effector_velocity=action)
-            penv.robot.action = jv
-            singularity = penv.robot.set_joint_velocities(penv.robot.action)
+            joint_vel, jacobian = robot.calculate_joint_velocities_from_ee_velocity_dls(end_effector_velocity=action)
+            robot.action = joint_vel
+            singularity = robot.set_joint_velocities(robot.action)
             penv.pbutils.pbclient.stepSimulation()
             time.sleep(0.01)
         except KeyboardInterrupt:
             break
 
-    # penv.deactivate_tree(tree_id_str="LPy_envy_tree1")
+    # penv.deactivate_tree(tree_id_str="LPy_tree1")
     penv.pbutils.pbclient.disconnect()
     return
 

pybullet_tree_sim/pruning_environment.py

@@ -55,57 +55,45 @@ class PruningEnv(gym.Env):
 
     def __init__(
         self,
-        # angle_threshold_perp: float = 0.52,
-        # angle_threshold_point: float = 0.52,
         pbutils: PyBUtils,
-        sensor_config: dict,
-        # cam_width: int = 424,
-        # cam_height: int = 240,
-        evaluate: bool = False,
-        # distance_threshold: float = 0.05,
-        max_steps: int = 1000,
+        # max_steps: int = 1000,
         make_trees: bool = False,
         name: str = "PruningEnv",
-        num_trees: int | None = None,
+        # num_trees: int | None = None,
         renders: bool = False,
         tree_count: int = 10,
         # tree_urdf_path: str | None = None,
         # tree_obj_path: str | None = None,
         verbose: bool = True,
-        load_robot: bool = True,
-        robot_type: str = "ur5",
-        robot_pos: ArrayLike = np.array([0, 0, 0]),
-        robot_orientation: ArrayLike = np.array([0, 0, 0, 1]),
+        # load_robot: bool = True,
+        robots: dict = {},
+        # robot_type: str = "ur5",
+        # robot_pos: ArrayLike = np.array([0, 0, 0]),
+        # robot_orientation: ArrayLike = np.array([0, 0, 0, 1]),
         # use_ik: bool = True,
         #
     ) -> None:
         """Initialize the Pruning Environment
 
         Parameters
         ----------
-            cam_width (int): Pixel width of the camera
-            cam_height (int): Pixel height of the camera
-            evaluate (bool): Is this environment for evaluation
-            max_steps (int): Maximum number of steps in a single run
-            name (str): Name of the environment (default="PruningEnv")
-            renders (bool): Whether to render the environment
-            tree_count (int): Number of trees to be loaded
+            ...
         """
         super().__init__()
         self.pbutils = pbutils
 
         # Pybullet GUI variables
         self.render_mode = "rgb_array"
         self.renders = renders
-        self.eval = evaluate
+        # self.eval = evaluate
 
         # Gym variables
         self.name = name
-        self.step_counter = 0
-        self.global_step_counter = 0
+        # self.step_counter = 0
+        # self.global_step_counter = 0
         # self.max_steps = max_steps
         self.tree_count = tree_count
-        self.is_goal_state = False
+        # self.is_goal_state = False
 
 
 
@@ -138,18 +126,19 @@ def __init__(
         self.debouce_time = 0.5
         self.last_button_push_time = time.time()
 
-        # Load Robot
-        if load_robot:
-            self.robot = Robot(pbclient=self.pbutils.pbclient)
+        # Load Robots
+        self.robots = robots
+        # if load_robot:
+        #     self.robot = Robot(pbclient=self.pbutils.pbclient)
             # self.ur5 = self.load_robot(
             #     type=robot_type, robot_pos=robot_pos, robot_orientation=robot_orientation, randomize_pose=False
             # )
-    
+            
         # Load all sensor attributes. # TODO: Load only the required sensor attributes
         self._load_sensor_attributes()
 
-        self.sensor_config = sensor_config
-        self._assign_tf_frame_to_sensors(self.sensor_config)
+        log.warn(self.robots['pruning_robot'].sensors)
+        # self.sensor_config = sensor_config
         # log.warning(self.sensor_attributes)
         return
 
@@ -169,13 +158,6 @@ def _load_sensor_attributes(self):
                 self.sensor_attributes[key] = value
         return
 
-    def _assign_tf_frame_to_sensors(self, sensor_config: dict):
-        for sensor_name, conf in sensor_config.items():
-            sensor = conf['sensor']
-            sensor.tf_frame = conf['tf_frame']
-            log.warn(f"{sensor.params}")
-            sensor.tf_frame_index = self.robot.robot_conf['joint_info']['mock_pruner__base--camera0']['id']
-        return
 
     def load_robot(self, type: str, robot_pos: ArrayLike, robot_orientation: ArrayLike, randomize_pose: bool = False):
         """Load a robot into the environment. Currently only UR5 is supported. TODO: Add Panda"""
@@ -468,7 +450,7 @@ def compute_deprojected_point_mask(self):
         # point_mask = np.expand_dims(point_mask_resize, axis=0).astype(np.float32)
         return point_mask
 
-    def is_reachable(self, vertex: Tuple[np.ndarray], base_xyz: np.ndarray) -> bool:
+    def is_reachable(self, robot: Robot, vertex: Tuple[np.ndarray], base_xyz: np.ndarray) -> bool:
         # if vertex[3] != "SPUR":
         #     return False
         ur5_base_pos = np.array(base_xyz)
@@ -479,7 +461,7 @@ def is_reachable(self, vertex: Tuple[np.ndarray], base_xyz: np.ndarray) -> bool:
         if dist >= 0.98:  # TODO: is this for the UR5? Should it be from a parameter file?
             return False
 
-        j_angles = self.robot.calculate_ik(vertex[0], None)
+        j_angles = robot.calculate_ik(vertex[0], None)
         # env.ur5.set_joint_angles(j_angles)
         # for _ in range(100):
         #     pyb.con.stepSimulation()
@@ -509,7 +491,7 @@ def get_key_pressed(self, relevant=None) -> list:
             keys_pressed.append(key)
         return keys_pressed
 
-    def get_key_action(self, keys_pressed: list) -> np.ndarray:
+    def get_key_action(self, robot: Robot, keys_pressed: list) -> np.ndarray:
         """Return an action based on the keys pressed."""
         action = np.array([0.0, 0.0, 0, 0.0, 0.0, 0.0])
         if keys_pressed:
@@ -540,9 +522,9 @@ def get_key_action(self, keys_pressed: list) -> np.ndarray:
             if ord("p") in keys_pressed:
                 if time.time() - self.last_button_push_time > self.debouce_time:
                     # Get view and projection matrices
-                    view_matrix = self.ur5.get_view_mat_at_curr_pose(pan=0, tilt=0, xyz_offset=[0, 0, 0])
+                    view_matrix = robot.get_view_mat_at_curr_pose(camera=robot.sensors['camera0']['sensor'])
                     log.warning(f"button p pressed")
-                    rgb, depth = self.pbutils.get_rgbd_at_cur_pose(type="robot", view_matrix=view_matrix)
+                    rgb, depth = self.pbutils.get_rgbd_at_cur_pose(camera=robot.sensors['camera0']['sensor'], type="robot", view_matrix=view_matrix)
                     # log.debug(f'depth:\n{depth}')
                     depth = -1 * depth.reshape((self.cam_width * self.cam_height, 1), order="F")
                     world_points = self.deproject_pixels_to_points(

pybullet_tree_sim/robot.py

@@ -22,6 +22,7 @@ class Robot:
     def __init__(
         self,
         pbclient,
+        sensor_config: dict = {},
         # robot_type: str,
         # robot_urdf_path: str,
         # tool_link_name: str,
@@ -32,7 +33,7 @@ def __init__(
         # robot_collision_filter_idxs,
         # init_joint_angles: Optional[list] = None,
         position=(0, 0, 0),
-        orientation=(0, 0, 0, 1),
+        orientation=(0, 0, 0, 1),   
         randomize_pose=False,
         verbose=1
         # **kwargs
@@ -58,8 +59,10 @@ def __init__(
         #     [0.063179, 0.077119, 0.0420027])
         self.verbose = verbose
 
-        self.joints = None
+        
         self.robot = None
+        self.joints = None
+        self.sensors = {}
         # self.robot_conf["control_joints"] = control_joints
         # self.robot_collision_filter_idxs = robot_collision_filter_idxs
         # 
@@ -79,6 +82,9 @@ def __init__(
         self._assign_control_joints()
         self.setup_robot()
 
+        self._assign_tf_frame_to_sensors(sensor_config)
+        log.warn(self.sensors)
+        
         self.action = None
 
         return
@@ -130,6 +136,14 @@ def _generate_robot_urdf(self) -> None:
         xutils.save_urdf(robot_urdf, urdf_path=self.robot_urdf_path)
 
         return
+        
+    def _assign_tf_frame_to_sensors(self, sensor_config: dict):
+        for sensor_name, conf in sensor_config.items():
+            sensor = conf['sensor']
+            sensor.tf_frame = conf['tf_frame']
+            log.warn(f"{sensor.params}")
+            sensor.tf_frame_index = self.robot_conf['joint_info']['mock_pruner__base--camera0']['id']
+        return
 
     def _assign_control_joints(self) -> None:
         self.robot_conf["control_joints"] = []
@@ -241,8 +255,8 @@ def setup_robot(self):
                     )
                 # self.joints[info.name] = info
 
-        import pprint as pp
-        pp.pprint(self.robot_conf)
+        # import pprint as pp
+        # pp.pprint(self.robot_conf)
 
         self.set_joint_angles_no_collision(self.init_joint_angles)
         self.pbclient.stepSimulation()
@@ -280,12 +294,17 @@ def _assign_collision_links(self):
         for i in range(num_joints):
             info = self.pbclient.getJointInfo(self.robot, i)
             child_link_name = info[12].decode('utf-8')
-
+            log.info(f"{child_link_name}, {i}")
             # This is kinda hacky, but it works for now. TODO: make better?
-            if child_link_name.endswith('base'):
-                self.robot_collision_filter_idxs.append((i, i-1))
-            # self.robot_conf["joint_info"].update({child_link_name: i})
-        log.warn(self.robot_collision_filter_idxs)
+        #     if child_link_name.endswith('base'):
+        #         self.robot_collision_filter_idxs.append((i, i-1))
+        #         log.warn(f'{child_link_name}, {i}')
+        #     if child_link_name.endswith('tool0'):
+        #         log.warn(f'{child_link_name}, {i}')
+        #     # self.robot_conf["joint_info"].update({child_link_name: i})
+        # self.robot_collision_filter_idxs.append((25, 2))
+
+        # log.warn(self.robot_collision_filter_idxs)
         # log.warn()
         return
 

pybullet_tree_sim/sensor.py

@@ -47,6 +47,9 @@ def main():
 
     sensor = Sensor(sensor_name="vl53l8cx", sensor_type="tof")
     print(sensor.params)
+    
+    sensor = Sensor(sensor_name="vl6180", sensor_type="tof")
+    print(sensor.params)
 
     return
 

pybullet_tree_sim/urdf/base_mounts/cart/macro/cart_macro.urdf.xacro

@@ -3,15 +3,11 @@
 
 <xacro:macro name="cart_macro" params="cart_prefix:='cart__' cart_parent:='world' mesh_base_path">
     <!-- ================== -->
-    <!--     Materials      -->
+    <!--     Cart Maacro    -->
     <!-- ================== -->
     <material name="LightGrey">
         <color rgba="0.7 0.7 0.7 1.0"/>
     </material>
-    
-    <!-- ================== -->
-    <!--     Properties     -->
-    <!-- ================== -->
     <xacro:property name="cart_mass" value="20.0"/>
     <xacro:property name="cart_collision_box_x" value='0.5'/>
     <xacro:property name="cart_collision_box_y" value='0.5'/>
@@ -24,14 +20,14 @@
 
     <link name="${cart_prefix}cart">
         <visual>
-            <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+            <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 ${cart_collision_box_z}"/>
             <geometry>
                 <mesh filename="${mesh_base_path}/base_mounts/cart/cart.obj"/>
             </geometry>
             <material name="LightGrey"/>
         </visual>
         <collision>
-            <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+            <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 -${cart_collision_box_z / 2}"/>
             <geometry>
                 <box size="${cart_collision_box_x} ${cart_collision_box_y} ${cart_collision_box_z}"/>
             </geometry>

pybullet_tree_sim/urdf/end_effectors/dovetail_mount/macro/dovetail_mount_macro.urdf.xacro

@@ -106,6 +106,9 @@
 
 
     <xacro:macro name="dovetail_mount_macro" params="parent_link dovetail_male_mount_prefix dovetail_female_mount_prefix mesh_base_path">
+        <!-- =================== -->
+        <!--   Dovetail Mount    -->
+        <!-- =================== -->
         <xacro:dovetail_male_mount_macro 
             parent_link="${parent_link}" 
             dovetail_male_mount_prefix="${dovetail_male_mount_prefix}"