set_robot_cell_state is working in PyBullet

docs/examples/05_backends_pybullet/files/01_set_robot_cell.py → docs/examples/05_backends_pybullet/files/01_robot_cell_basics.py

@@ -7,13 +7,15 @@
 from compas_fab.robots import RigidBody
 
 with PyBulletClient() as client:
-    urdf_filepath = compas_fab.get("robot_library/ur5_robot/urdf/robot_description.urdf")
-    robot = client.load_robot(urdf_filepath)
 
-    robot_cell = RobotCell(robot)
+    # Create a robot cell and add objects to it
+    robot_cell = RobotCell()  # Typically RobotCell is initialed with a RobotModel, but here it is empty for simplicity
     floor_mesh = Mesh.from_stl(compas_fab.get("planning_scene/floor.stl"))
     robot_cell.rigid_body_models["floor"] = RigidBody(floor_mesh)
+    cone = Mesh.from_stl(compas_fab.get("planning_scene/cone.stl"))
+    robot_cell.rigid_body_models["cone"] = RigidBody(cone)
 
+    # The planner is used for passing the robot cell into the PyBullet client
     planner = PyBulletPlanner(client)
     planner.set_robot_cell(robot_cell)
     # The floor should appear in the PyBullet's GUI

docs/examples/05_backends_pybullet/files/01_robot_cell_with_tools.py

@@ -0,0 +1,69 @@
+from compas.datastructures import Mesh
+
+import compas_fab
+from compas_fab.backends import PyBulletClient
+from compas_fab.backends import PyBulletPlanner
+from compas_fab.robots import RobotCell
+from compas_fab.robots import RobotCellState
+from compas_fab.robots import RigidBody
+from compas_fab.robots import RobotLibrary
+from compas.geometry import Frame
+from compas.geometry import Box
+from compas_robots import ToolModel
+
+with PyBulletClient() as client:
+    robot = RobotLibrary.abb_irb4600_40_255()
+    robot = client.load_existing_robot(robot)
+
+    # ---------------------------------------------------------------------
+    # Create a robot cell and add objects to it
+    # ---------------------------------------------------------------------
+    robot_cell = RobotCell(robot)
+
+    # Add Static Collision Geometry
+    floor_mesh = Mesh.from_stl(compas_fab.get("planning_scene/floor.stl"))
+    robot_cell.rigid_body_models["floor"] = RigidBody(floor_mesh)
+
+    # Add Tool
+    tool_mesh = Mesh.from_stl(compas_fab.get("planning_scene/cone.stl"))
+    tool_frame = Frame([0, 0, 0.14], [1, 0, 0], [0, 1, 0])
+    robot_cell.tool_models["cone"] = ToolModel(tool_mesh, tool_frame)
+
+    # Add workpiece at tool tip
+    workpiece_mesh = Box(1.0, 0.1, 0.2).to_mesh(triangulated=True)
+    robot_cell.rigid_body_models["workpiece"] = RigidBody(workpiece_mesh)
+
+    # ------------------------------------------------------------------------
+    # Create a RobotCellState to represent the current state of the robot cell
+    # ------------------------------------------------------------------------
+    robot_cell_state = RobotCellState.from_robot_cell(robot_cell)
+
+    # Change the robot's configuration for demonstration purposes
+    configuration = robot.zero_configuration()
+    configuration.joint_values[1] = 0.5  # Change the second joint angle to 0.5 [rad]
+    robot_cell_state.robot_configuration = configuration
+
+    # Attach the tool to the robot's main group
+    robot_cell_state.set_tool_attached_to_group("cone", robot.main_group_name)
+
+    # Attach the workpiece to the tool
+    workpiece_grasp_frame = Frame([0, 0, 0.1], [1, 0, 0], [0, 1, 0])
+    robot_cell_state.set_rigid_body_attached_to_tool("workpiece", "cone", workpiece_grasp_frame)
+
+    # The planner is used for passing the robot cell into the PyBullet client
+    planner = PyBulletPlanner(client)
+    planner.set_robot_cell(robot_cell)  # or planner.set_robot_cell(robot_cell, robot_cell_state)
+    planner.set_robot_cell_state(robot_cell_state)
+
+    # Typically the robot_cell_state is passed to the
+    # planning functions such as planner.plan_motion(start_state, target), or
+    # visualization functions such as robot_cell_scene_object.update(robot_cell_state).
+    # In this example, we are directly calling set_robot_cell_state() to see the effect,
+    # which can be seen in the PyBullet GUI.
+
+    input("Observe the PyBullet GUI, Press Enter to continue...")
+
+    for i in range(10):
+        robot_cell_state.robot_configuration.joint_values[1] += 0.1
+        planner.set_robot_cell_state(robot_cell_state)
+        input("Observe the PyBullet GUI, Press Enter to continue...")

src/compas_fab/backends/interfaces/backend_features.py

@@ -33,18 +33,24 @@ class BackendFeature(object):
     """
 
     def __init__(self, client=None):
-        # All backend features are assumed to be associated with a backend client.
-        if client:
-            self.client = client  # type: ClientInterface
-
-        # The current robot cell last set by user calling Planner.set_robot_cell()
-        self.robot_cell = None  # type: RobotCell
-
-        # The current robot cell state last set by user calling Planner.set_robot_cell_state()
-        self.robot_cell_state = None  # type: RobotCellState
-
         super(BackendFeature, self).__init__()
 
+    # @property
+    # def client(self):
+    #     # type: () -> ClientInterface
+    #     """Proxy function to access the backend client.
+    #     This function should be overridden by the PlannerInterface default :meth:`PlannerInterface.client` or by the Planner.
+    #     """
+    #     raise NotImplementedError
+
+    # @property
+    # def robot_cell(self):
+    #     # type: () -> RobotCell
+    #     """Proxy function to access the RobotCell object.
+    #     This function should be overridden by the PlannerInterface default :meth:`PlannerInterface.robot_cell` or by the Planner.
+    #     """
+    #     raise NotImplementedError
+
 
 #   The code that contains the actual feature implementation is located in the backend's module.
 #   For example, the features for moveit planner and ros client are located in :

src/compas_fab/backends/pybullet/backend_features/__init__.py

@@ -33,6 +33,7 @@
 )
 from compas_fab.backends.pybullet.backend_features.pybullet_remove_collision_mesh import PyBulletRemoveCollisionMesh
 from compas_fab.backends.pybullet.backend_features.pybullet_set_robot_cell import PyBulletSetRobotCell
+from compas_fab.backends.pybullet.backend_features.pybullet_set_robot_cell_state import PyBulletSetRobotCellState
 
 
 __all__ = [
@@ -44,4 +45,5 @@
     "PyBulletRemoveAttachedCollisionMesh",
     "PyBulletRemoveCollisionMesh",
     "PyBulletSetRobotCell",
+    "PyBulletSetRobotCellState",
 ]

src/compas_fab/backends/pybullet/backend_features/pybullet_forward_kinematics.py

@@ -4,6 +4,14 @@
 
 from compas_fab.backends.interfaces import ForwardKinematics
 
+import compas
+
+if compas.IPY:
+    from typing import TYPE_CHECKING
+
+    if TYPE_CHECKING:
+        from compas_fab.backends import PyBulletClient  # noqa: F401
+
 
 class PyBulletForwardKinematics(ForwardKinematics):
     """Callable to calculate the robot's forward kinematic."""
@@ -38,12 +46,14 @@ def forward_kinematics(self, robot, configuration, group=None, options=None):
         """
         options = options or {"link": None, "check_collision": False}
 
+        client = self.client  # type: PyBulletClient # Trick to keep intellisense happy
+
         link_name = options.get("link") or robot.get_end_effector_link_name(group)
-        cached_robot_model = self.client.get_cached_robot_model(robot)
-        body_id = self.client.get_uid(cached_robot_model)
-        link_id = self.client._get_link_id_by_name(link_name, cached_robot_model)
-        self.client.set_robot_configuration(robot, configuration, group)
-        frame = self.client._get_link_frame(link_id, body_id)
+        cached_robot_model = client.get_cached_robot_model(robot)
+        body_id = client.get_uid(cached_robot_model)
+        link_id = client._get_link_id_by_name(link_name, cached_robot_model)
+        client.set_robot_configuration(robot, configuration, group)
+        frame = client._get_link_frame(link_id, body_id)
         if options.get("check_collision"):
-            self.client.check_collisions(robot, configuration)
+            client.check_collisions(robot, configuration)
         return frame

src/compas_fab/backends/pybullet/backend_features/pybullet_inverse_kinematics.py

@@ -102,7 +102,7 @@ def iter_inverse_kinematics(self, target, start_state=None, group=None, options=
         if robot.need_scaling:
             target_frame = target_frame.scaled(1.0 / robot.scale_factor)
             # Now target_frame is back in meter scale
-        attached_tool = self.client.robot_cell.get_attached_tool(start_state, group)
+        attached_tool = self.robot_cell.get_attached_tool(start_state, group)
         if attached_tool:
             target_frame = from_tcf_to_t0cf(target_frame, attached_tool.frame)
             # Attached tool frames does not need scaling because Tools are modelled in meter scale

src/compas_fab/backends/pybullet/backend_features/pybullet_set_robot_cell.py

@@ -1,6 +1,7 @@
 from compas_fab.backends.interfaces import SetRobotCell
 
 import compas
+from compas.geometry import Frame
 
 if not compas.IPY:
     from typing import TYPE_CHECKING
@@ -42,18 +43,18 @@ def set_robot_cell(self, robot_cell, robot_cell_state=None, options=None):
         # and the new robot cell is added to the PyBullet world
 
         # Remove all objects from the PyBullet world
-        client.remove_all_tools()
-        client.remove_all_rigid_bodies()
+        # client.remove_all_tools()
+        # client.remove_all_rigid_bodies()
 
         # Add the robot cell to the PyBullet world
-        for name, tool in robot_cell.tool_models.items():
-            client.add_tool(name, tool.mesh, tool.mass)
+        for name, tool_model in robot_cell.tool_models.items():
+            client.add_tool(name, tool_model)
         for name, rigid_body in robot_cell.rigid_body_models.items():
-            client.add_rigid_body(name, rigid_body.mesh, rigid_body.mass)
-
+            client.add_rigid_body(name, rigid_body)
+            # client.convert_mesh_to_body(rigid_body.visual_meshes[0], Frame.worldXY())
         # Update the robot cell in the client
         self._robot_cell = robot_cell
 
         # If a robot cell state is provided, update the client's robot cell state
-        if robot_cell_state:
-            self.set_robot_cell_state(robot_cell_state, options)
+        # if robot_cell_state:
+        #     self.set_robot_cell_state(robot_cell_state, options)

src/compas_fab/backends/pybullet/backend_features/pybullet_set_robot_cell_state.py

@@ -0,0 +1,90 @@
+from compas_fab.backends.interfaces import SetRobotCellState
+
+import compas
+from compas.geometry import Frame
+from compas.geometry import Transformation
+
+if not compas.IPY:
+    from typing import TYPE_CHECKING
+
+    if TYPE_CHECKING:
+        from typing import Optional  # noqa: F401
+        from typing import Dict  # noqa: F401
+        from typing import List  # noqa: F401
+        from typing import Tuple  # noqa: F401
+
+        from compas_fab.robots import Robot  # noqa: F401
+        from compas_robots import Configuration  # noqa: F401
+        from compas.geometry import Frame  # noqa: F401
+        from compas_fab.backends.interfaces import ClientInterface  # noqa: F401
+        from compas_fab.robots import RobotCell  # noqa: F401
+        from compas_fab.robots import RobotCellState  # noqa: F401
+        from compas_fab.backends import PyBulletClient  # noqa: F401
+
+from compas_fab.backends.pybullet.const import STATIC_MASS
+
+
+class PyBulletSetRobotCellState(SetRobotCellState):
+
+    def set_robot_cell_state(self, robot_cell_state):
+        # type: (RobotCellState, RobotCellState) -> None
+        """Change the state of the models in the robot cell that have already been set to the Pybullet client.
+
+        This function is typically called by planning functions that accepts a start state as input.
+        The user should not need to call this function directly.
+
+        The robot cell state must match the robot cell set earlier by :meth:`set_robot_cell`.
+
+        """
+        client = self.client  # type: PyBulletClient # Trick to keep intellisense happy
+        robot_cell = self.robot_cell  # type: RobotCell
+
+        robot_cell.assert_cell_state_match(robot_cell_state)
+
+        # Update the robot configuration
+        # Note robot_cell_state.robot_configuration is a full configuration
+        client.set_robot_configuration(robot_cell_state.robot_configuration)
+
+        # Keep track of tool's base_frames
+        tool_base_frames = {}
+
+        # Update the position of Tool models
+        for tool_name, tool_state in robot_cell_state.tool_states.items():
+            # if tool_state.is_hidden:
+            #     client.hide_tool(tool_name)
+            #     continue
+            if tool_state.attached_to_group:
+                # If tool is attached to a group, update the tool's base frame using the group's FK frame
+                tool_attached_link_name = client.robot.get_link_names(tool_state.attached_to_group)[-1]
+                configuration = robot_cell_state.robot_configuration
+                tool_base_frame = client.robot.model.forward_kinematics(configuration, tool_attached_link_name)
+            else:
+                # If the tool is not attached, update the tool's base frame using the frame in tool state
+                tool_base_frame = tool_state.frame
+            tool_base_frames[tool_name] = tool_base_frame
+            client.set_tool_base_frame(tool_name, tool_base_frame)
+
+        # Update the position of RigidBody models
+        for rigid_body_name, rigid_body_state in robot_cell_state.rigid_body_states.items():
+            # if rigid_body_state.is_hidden:
+            #     client.hide_rigid_body(rigid_body_name)
+            #     continue
+            if rigid_body_state.attached_to_tool:
+                # If rigid body is attached to a tool, update the rigid body's base frame using the tool's FK frame
+                tool_name = rigid_body_state.attached_to_tool
+                tool_model = robot_cell.tool_models[tool_name]
+
+                # t_tcf_ocf is Grasp, describing Object Coordinate Frame (OCF) relative to Tool Coordinate Frame (TCF)
+                t_tcf_ocf = Transformation.from_frame(rigid_body_state.grasp)
+                # t_t0cf_tcf is Tool Offset, describing Tool Coordinate Frame (TCF) relative to Tool Base Frame (T0CF)
+                t_t0cf_tcf = Transformation.from_frame(tool_model.frame)
+                # t0cf_from_wcf is Tool Base Frame, describing Tool Base Frame (T0CF) relative to World Coordinate Frame (WCF)
+                t_wcf_t0cf = Transformation.from_frame(tool_base_frames[tool_name])
+                # ocf_from_wcf is Object Coordinate Frame (OCF) relative to World Coordinate Frame (WCF)
+                t_wcf_ocf = t_wcf_t0cf * t_t0cf_tcf * t_tcf_ocf
+
+                rigid_body_base_frame = Frame.from_transformation(t_wcf_ocf)
+            else:
+                # If the rigid body is not attached, update the rigid body's base frame using the frame in rigid body state
+                rigid_body_base_frame = rigid_body_state.frame
+            client.set_rigid_body_base_frame(rigid_body_name, rigid_body_base_frame)