Add whole-body differential IK example (#66)

docs/source/motion_planning.rst

@@ -87,7 +87,7 @@ However, for most robotics applications, we rely on **numerical methods** instea
 * Additional constraints, such as joint limits, Cartesian pose limits, or collision avoidance, are difficult to enforce analytically.
 
 The `pyroboplan.ik module <api/pyroboplan.ik.html>`_ contains implementations for IK solvers.
-You can also try running the :examples:`differential IK example <differential_ik.py>`.
+You can also try running the :examples:`fixed arm differential IK <differential_ik.py>` and :examples:`whole-body differential IK <differential_ik_whole_body.py>` examples.
 
 .. image:: _static/images/inverse_kinematics.png
     :width: 600

examples/differential_ik.py

@@ -78,7 +78,10 @@
     for _ in range(10):
         init_state = get_random_collision_free_state(model, collision_model)
         target_tform = get_random_collision_free_transform(
-            model, collision_model, target_frame
+            model,
+            collision_model,
+            target_frame,
+            joint_padding=0.05,
         )
         q_sol = ik.solve(
             target_frame,

examples/differential_ik_whole_body.py

@@ -0,0 +1,92 @@
+"""
+This example shows PyRoboPlan capabilities for whole-body inverse kinematics (IK),
+using a 6-DOF arm on a mobile base capable of planar (3-DOF) motion.
+
+IK defines the task of finding a set of joint positions for a robot model to
+achieve a desired target pose for a specific coordinate frame.
+"""
+
+from pinocchio.visualize import MeshcatVisualizer
+import numpy as np
+
+from pyroboplan.core.utils import (
+    get_random_collision_free_state,
+    get_random_collision_free_transform,
+)
+from pyroboplan.ik.differential_ik import DifferentialIk, DifferentialIkOptions
+from pyroboplan.ik.nullspace_components import (
+    joint_limit_nullspace_component,
+    collision_avoidance_nullspace_component,
+)
+from pyroboplan.models.ur5 import (
+    load_ur5_on_base_models,
+    add_ur5_on_base_self_collisions,
+)
+
+
+if __name__ == "__main__":
+    # Create models and data
+    model, collision_model, visual_model = load_ur5_on_base_models()
+    add_ur5_on_base_self_collisions(model, collision_model)
+
+    data = model.createData()
+    collision_data = collision_model.createData()
+
+    target_frame = "tool0"
+
+    # Initialize visualizer
+    viz = MeshcatVisualizer(model, collision_model, visual_model, data=data)
+    viz.initViewer(open=True)
+    viz.loadViewerModel()
+    np.set_printoptions(precision=3)
+
+    # Set up the IK solver.
+    # The first 3 values in the weights are the base, and the other 6 are the arm.
+    # This means we are putting a higher cost on base motion more than arm motion.
+    options = DifferentialIkOptions(
+        damping=0.0001,
+        min_step_size=0.025,
+        max_step_size=0.1,
+        joint_weights=[5.0, 5.0, 15.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+        max_retries=5,
+        rng_seed=None,
+    )
+    ik = DifferentialIk(
+        model,
+        data=data,
+        collision_model=collision_model,
+        options=options,
+        visualizer=viz,
+    )
+    nullspace_components = [
+        lambda model, q: collision_avoidance_nullspace_component(
+            model,
+            data,
+            collision_model,
+            collision_data,
+            q,
+            gain=0.25,
+            dist_padding=0.02,
+        ),
+        lambda model, q: joint_limit_nullspace_component(
+            model, q, gain=0.1, padding=0.02
+        ),
+    ]
+
+    # Solve IK several times and print the results
+    for _ in range(10):
+        init_state = get_random_collision_free_state(model, collision_model)
+        target_tform = get_random_collision_free_transform(
+            model,
+            collision_model,
+            target_frame,
+            joint_padding=0.05,
+        )
+        q_sol = ik.solve(
+            target_frame,
+            target_tform,
+            init_state=init_state,
+            nullspace_components=nullspace_components,
+            verbose=True,
+        )
+        print(f"Solution configuration:\n{q_sol}\n")

examples/intro_pinocchio/intro_pinocchio_ur5_sim.py

@@ -4,28 +4,29 @@
 Jacobian computation, collision checking, and dynamics simulation.
 """
 
-import pinocchio
-from pinocchio.visualize import MeshcatVisualizer
-
 import meshcat.geometry as mg
 import numpy as np
-from os.path import dirname, join, abspath
+import os
 import time
 
+import pinocchio
+from pinocchio.visualize import MeshcatVisualizer
+from pyroboplan.models.utils import get_example_models_folder
+
 
 def main():
     # Load the models from the URDF file.
-    pinocchio_model_dir = join(dirname(str(abspath(__file__))), "..", "..", "models")
-    urdf_filename = join(
-        pinocchio_model_dir, "ur_description", "urdf", "ur5_gripper.urdf"
+    model_dir = get_example_models_folder()
+    urdf_filename = os.path.join(
+        model_dir, "ur5_description", "urdf", "ur5_gripper.urdf"
     )
     model, collision_model, visual_model = pinocchio.buildModelsFromUrdf(
-        urdf_filename, package_dirs=pinocchio_model_dir
+        urdf_filename, package_dirs=model_dir
     )
 
     # Modify the collision model for display.
-    srdf_filename = join(
-        pinocchio_model_dir, "ur_description", "srdf", "ur5_gripper.srdf"
+    srdf_filename = os.path.join(
+        model_dir, "ur5_description", "srdf", "ur5_gripper.srdf"
     )
     collision_model.addAllCollisionPairs()
     pinocchio.removeCollisionPairs(model, collision_model, srdf_filename)