Add kinematic trajectory optimization (sea-bass#37)

docs/source/motion_planning.rst

@@ -115,7 +115,7 @@ Currently, all the planners in ``pyroboplan`` (such as RRT and Cartesian interpo
 Online planning and control is often done through optimization techniques like Model Predictive Control (MPC).
 
 The `pyroboplan.planning module <api/pyroboplan.planning.html>`_ contains implementations for a number of motion planners.
-You can also try running the :examples:`RRT example <example_rrt.py>`, :examples:`PRM example <example_prm.py>`, or the :examples:`Cartesian planning example <example_cartesian_path.py>`.
+You can also try running the :examples:`RRT example <example_rrt.py>`, :examples:`PRM example <example_prm.py>`, or :examples:`Cartesian planning example <example_cartesian_path.py>`.
 
 .. image:: _static/images/bidirectional_rrt_star.png
     :width: 600
@@ -137,8 +137,8 @@ As mentioned in the previous section, if you are using a planner that simply out
 Often, a fixed set of kinematic (position/velocity/acceleration/jerk) and dynamic (force/torque) limits of the robot are taken into account.
 Sometimes, these limits can also be task-dependent; for example, if manipulating fragile objects or objects that cannot be placed in certain configurations (e.g., moving a glass of water without spilling).
 
-The `pyroboplan.trajectory module <api/pyroboplan.trajectory.html>`_ contains trajectory generation implementations.
-You can also try running the :examples:`trajectory generation example <example_trajectory.py>`.
+The `pyroboplan.trajectory module <api/pyroboplan.trajectory.html>`_ contains trajectory generation and optimization implementations.
+You can try running the corresponding :examples:`trajectory generation <example_trajectory_generation.py>` and :examples:`trajectory optimization <example_trajectory_optimization.py>` examples.
 
 .. image:: _static/images/trajectory_generation.png
     :width: 720

examples/example_cartesian_path.py

@@ -83,7 +83,7 @@
     plt.plot(t_vec, q_vec[idx, :])
 
 curr_time = 0
-time_line = plt.axvline(x=curr_time, color="b")
+time_line = plt.axvline(x=curr_time, color="k", linestyle="--")
 
 plt.legend(model.names[1:])
 plt.show()

examples/example_trajectory.py → examples/example_trajectory_generation.py

@@ -47,17 +47,24 @@
         qd_max = 0.75
         qdd_max = 0.5
         traj = TrapezoidalVelocityTrajectory(q, qd_max, qdd_max)
+        t_vec, q_vec, qd_vec, qdd_vec = traj.generate(dt)
     elif mode == "quintic":
         t_vec = [0.0, 3.0, 6.0]
         qd = 0.0
         qdd = 0.0
         traj = QuinticPolynomialTrajectory(t_vec, q, qd, qdd)
-
-    t_vec, q_vec, qd_vec, qdd_vec = traj.generate(dt)
+        t_vec, q_vec, qd_vec, qdd_vec, qddd_vec = traj.generate(dt)
 
     # Display the trajectory and points along the path.
+    # If using a polynomial trajectory, you can also add show_jerk=True.
     plt.ion()
-    traj.visualize(dt=dt, joint_names=model.names[1:])
+    traj.visualize(
+        dt=dt,
+        joint_names=model.names[1:],
+        show_position=True,
+        show_velocity=True,
+        show_acceleration=True,
+    )
     time.sleep(0.5)
 
     tforms = extract_cartesian_poses(model, "panda_hand", [q_start, q_mid, q_end])

examples/example_trajectory_optimization.py

@@ -0,0 +1,91 @@
+"""
+This example shows PyRoboPlan capabilities for path planning using
+trajectory optimization.
+"""
+
+import matplotlib.pyplot as plt
+from pinocchio.visualize import MeshcatVisualizer
+import time
+
+from pyroboplan.core.utils import (
+    get_random_collision_free_state,
+    extract_cartesian_poses,
+)
+from pyroboplan.models.panda import load_models, add_self_collisions
+from pyroboplan.trajectory.trajectory_optimization import (
+    CubicTrajectoryOptimization,
+    CubicTrajectoryOptimizationOptions,
+)
+from pyroboplan.visualization.meshcat_utils import visualize_frames
+
+
+if __name__ == "__main__":
+    # Create models and data.
+    model, collision_model, visual_model = load_models()
+    add_self_collisions(model, collision_model)
+    data = model.createData()
+
+    # Initialize visualizer.
+    viz = MeshcatVisualizer(model, collision_model, visual_model, data=data)
+    viz.initViewer(open=True)
+    viz.loadViewerModel()
+    q_start = get_random_collision_free_state(model, collision_model)
+    viz.display(q_start)
+    time.sleep(1.0)
+
+    # Configure trajectory optimization.
+    dt = 0.025
+    options = CubicTrajectoryOptimizationOptions(
+        num_waypoints=5,
+        samples_per_segment=11,
+        min_segment_time=0.01,
+        max_segment_time=10.0,
+        min_vel=-1.5,
+        max_vel=1.5,
+        min_accel=-0.75,
+        max_accel=0.75,
+    )
+
+    # Perform trajectory optimization.
+    multi_point = True
+    if multi_point:
+        # Multi point means we set all the waypoints and optimize how to move between them.
+        q_path = [q_start] + [
+            get_random_collision_free_state(model, collision_model)
+            for _ in range(options.num_waypoints - 1)
+        ]
+    else:
+        # Single point means we set just the start and goal.
+        # All other intermediate waypoints are optimized automatically.
+        q_path = [q_start, get_random_collision_free_state(model, collision_model)]
+
+    planner = CubicTrajectoryOptimization(model, options)
+    traj = planner.plan(q_path)
+
+    if traj is not None:
+        print("Trajectory optimization successful")
+        traj_gen = traj.generate(dt)
+        q_vec = traj_gen[1]
+
+        # Display the trajectory and points along the path.
+        plt.ion()
+        traj.visualize(
+            dt=dt,
+            joint_names=model.names[1:],
+            show_position=True,
+            show_velocity=True,
+            show_acceleration=True,
+            show_jerk=True,
+        )
+        time.sleep(0.5)
+
+        tforms = extract_cartesian_poses(model, "panda_hand", q_vec.T)
+        viz.display(q_start)
+        visualize_frames(viz, "waypoints", tforms, line_length=0.075, line_width=2)
+        time.sleep(1.0)
+
+        # Animate the generated trajectory.
+        input("Press 'Enter' to animate the path.")
+        for idx in range(q_vec.shape[1]):
+            viz.display(q_vec[:, idx])
+            time.sleep(dt)

pyproject.toml

@@ -6,6 +6,7 @@ build-backend = "hatchling.build"
 name = "pyroboplan"
 version = "1.0.0"
 dependencies = [
+  "drake == 1.28.0",
   "pin == 2.7.0",
   "matplotlib == 3.8.4",
   "meshcat == 0.3.2",

src/pyroboplan/planning/rrt.py

@@ -119,7 +119,7 @@ def plan(self, q_start, q_goal):
         ----------
             q_start : array-like
                 The starting robot configuration.
-            q_start : array-like
+            q_goal : array-like
                 The goal robot configuration.
         """
         self.reset()

src/pyroboplan/trajectory/polynomial.py

@@ -25,13 +25,14 @@ def evaluate(self, t):
 
         Returns
         -------
-            tuple(array-like or float, array-like or float, array-like or float)
-                The tuple of (q, qd, qdd) trajectory values at the specified time.
+            tuple(array-like or float, array-like or float, array-like or float, array-like or float)
+                The tuple of (q, qd, qdd, qddd) trajectory values at the specified time.
                 If the trajectory is single-DOF, these will be scalars, otherwise they are arrays.
         """
-        q = np.zeros(self.num_dims)
-        qd = np.zeros(self.num_dims)
-        qdd = np.zeros(self.num_dims)
+        q = np.zeros(self.num_dims)  # Position
+        qd = np.zeros(self.num_dims)  # Velocity
+        qdd = np.zeros(self.num_dims)  # Acceleration
+        qddd = np.zeros(self.num_dims)  # Jerk
 
         # Handle edge cases.
         if t < self.segment_times[0][0]:
@@ -53,6 +54,7 @@ def evaluate(self, t):
                     q[dim] = np.polyval(coeffs, dt)
                     qd[dim] = np.polyval(np.polyder(coeffs, 1), dt)
                     qdd[dim] = np.polyval(np.polyder(coeffs, 2), dt)
+                    qddd[dim] = np.polyval(np.polyder(coeffs, 3), dt)
                 evaluated = True
             else:
                 segment_idx += 1
@@ -62,7 +64,8 @@ def evaluate(self, t):
             q = q[0]
             qd = qd[0]
             qdd = qdd[0]
-        return q, qd, qdd
+            qddd = qddd[0]
+        return q, qd, qdd, qddd
 
     def generate(self, dt):
         """
@@ -75,8 +78,8 @@ def generate(self, dt):
 
         Returns
         -------
-            tuple(array-like, array-like, array-like, array-like)
-                The tuple of (t, q, qd, qdd) trajectory values generated at the sample time.
+            tuple(array-like, array-like, array-like, array-like, array-like)
+                The tuple of (t, q, qd, qdd, qddd) trajectory values generated at the sample time.
                 The time vector is 1D, and the others are 2D (time and dimension).
         """
 
@@ -88,9 +91,10 @@ def generate(self, dt):
             t_vec = np.append(t_vec, t_final)
         num_pts = len(t_vec)
 
-        q = np.zeros([self.num_dims, num_pts])
-        qd = np.zeros([self.num_dims, num_pts])
-        qdd = np.zeros([self.num_dims, num_pts])
+        q = np.zeros([self.num_dims, num_pts])  # Position
+        qd = np.zeros([self.num_dims, num_pts])  # Velocity
+        qdd = np.zeros([self.num_dims, num_pts])  # Acceleration
+        qddd = np.zeros([self.num_dims, num_pts])  # Jerk
 
         t_idx = 0
         segment_idx = 0
@@ -106,13 +110,22 @@ def generate(self, dt):
                     q[dim, t_idx] = np.polyval(coeffs, dt)
                     qd[dim, t_idx] = np.polyval(np.polyder(coeffs, 1), dt)
                     qdd[dim, t_idx] = np.polyval(np.polyder(coeffs, 2), dt)
+                    qddd[dim, t_idx] = np.polyval(np.polyder(coeffs, 3), dt)
                 t_idx += 1
             else:
                 segment_idx += 1
 
-        return t_vec, q, qd, qdd
-
-    def visualize(self, dt=0.01, joint_names=None):
+        return t_vec, q, qd, qdd, qddd
+
+    def visualize(
+        self,
+        dt=0.01,
+        joint_names=None,
+        show_position=True,
+        show_velocity=False,
+        show_acceleration=False,
+        show_jerk=False,
+    ):
         """
         Visualizes a generated trajectory with one figure window per dimension.
 
@@ -123,10 +136,18 @@ def visualize(self, dt=0.01, joint_names=None):
             joint_names : list[str], optional
                 The joint names to use for the figure titles.
                 If unset, uses the text "Dimension 1", "Dimension 2", etc.
+            show_position : bool, optional
+                If true, shows the trajectory position values.
+            show_velocity : bool, optional
+                If true, shows the trajectory velocity values.
+            show_acceleration : bool, optional
+                If true, shows the trajectory acceleration values.
+            show_jerk : bool, optional
+                If true, shows the trajectory jerk values.
         """
         import matplotlib.pyplot as plt
 
-        t_vec, q, qd, qdd = self.generate(dt)
+        t_vec, q, qd, qdd, qddd = self.generate(dt)
         vertical_line_scale_factor = 1.2
 
         for dim in range(self.num_dims):
@@ -137,21 +158,37 @@ def visualize(self, dt=0.01, joint_names=None):
             plt.figure(title)
 
             # Positions, velocities, and accelerations
-            plt.plot(t_vec, q[dim, :])
-            plt.plot(t_vec, qd[dim, :])
-            plt.plot(t_vec, qdd[dim, :])
-            plt.legend(["Position", "Velocity", "Acceleration"])
+            legend = []
+            min_pos = min_vel = min_accel = min_jerk = np.inf
+            max_pos = max_vel = max_accel = max_jerk = -np.inf
+            if show_position:
+                plt.plot(t_vec, q[dim, :])
+                min_pos = np.min(q[dim, :])
+                max_pos = np.max(q[dim, :])
+                legend.append("Position")
+            if show_velocity:
+                plt.plot(t_vec, qd[dim, :])
+                min_vel = np.min(qd[dim, :])
+                max_vel = np.max(qd[dim, :])
+                legend.append("Velocity")
+            if show_acceleration:
+                plt.plot(t_vec, qdd[dim, :])
+                min_accel = np.min(qdd[dim, :])
+                max_accel = np.max(qdd[dim, :])
+                legend.append("Acceleration")
+            if show_jerk:
+                plt.plot(t_vec, qddd[dim, :])
+                min_jerk = np.min(qddd[dim, :])
+                max_jerk = np.max(qddd[dim, :])
+                legend.append("Jerk")
+            plt.legend(legend)
 
             # Times
             min_val = vertical_line_scale_factor * min(
-                np.min(q[dim, :]),
-                np.min(qd[dim, :]),
-                np.min(qdd[dim, :]),
+                [min_pos, min_vel, min_accel, min_jerk]
             )
             max_val = vertical_line_scale_factor * max(
-                np.max(q[dim, :]),
-                np.max(qd[dim, :]),
-                np.max(qdd[dim, :]),
+                [max_pos, max_vel, max_accel, max_jerk]
             )
             plt.vlines(
                 [t[1] for t in self.segment_times],