Fix dual Cartesian DMP in Cython

dfki-ric · Nov 18, 2023 · cff7356 · cff7356
1 parent 2a62e43
commit cff7356
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 15 deletions.
diff --git a/examples/plot_dmp_scaling.py b/examples/plot_dmp_scaling.py
@@ -24,19 +24,19 @@
 y = np.cos(T ** 2 * 1.99 * np.pi)
 z = qx = qy = qz = np.zeros_like(x)
 qw = np.ones_like(x)
-Y = np.column_stack((x, y, z, qw, qx, qy, qz))#, x, y, z, qw, qx, qy, qz))
+Y = np.column_stack((x, y, z, qw, qx, qy, qz, x, y, z, qw, qx, qy, qz))
 start = Y[0]
 goal = Y[-1]
-new_start = np.array([0.5, 0.5, 0, 1, 0, 0, 0])#, 0.5, 0.5, 0, 1, 0, 0, 0])
-new_goal = np.array([-0.5, 0.5, 0, 1, 0, 0, 0])#, -0.5, 0.5, 0, 1, 0, 0, 0])
+new_start = np.array([0.5, 0.5, 0, 1, 0, 0, 0, 0.5, 0.5, 0, 1, 0, 0, 0])
+new_goal = np.array([-0.5, 0.5, 0, 1, 0, 0, 0, -0.5, 0.5, 0, 1, 0, 0, 0])
 Y_shifted = Y - goal[np.newaxis] + new_goal[np.newaxis]
 
-dmp = DMP(n_dims=len(start), execution_time=1.0, dt=0.01, n_weights_per_dim=20)
+#dmp = DMP(n_dims=len(start), execution_time=1.0, dt=0.01, n_weights_per_dim=20)
 #dmp = CartesianDMP(execution_time=1.0, dt=0.01, n_weights_per_dim=20)
-#dmp = DualCartesianDMP(execution_time=1.0, dt=0.01, n_weights_per_dim=20)
+dmp = DualCartesianDMP(execution_time=1.0, dt=0.01, n_weights_per_dim=20)
 dmp.imitate(T, Y)
 dmp.configure(start_y=new_start, goal_y=new_goal)
-_, Y_dmp = dmp.open_loop(step_function="rk4-cython")
+_, Y_dmp = dmp.open_loop(step_function="cython")
 
 plt.plot(Y[:, 0], Y[:, 1], label=r"Demonstration, $g \approx y_0$", ls="--")
 plt.plot(Y_shifted[:, 0], Y_shifted[:, 1], label="Original shape with new goal", ls="--")

diff --git a/movement_primitives/dmp_fast.pyx b/movement_primitives/dmp_fast.pyx
@@ -573,6 +573,7 @@ cpdef dmp_step_dual_cartesian(
     cdef np.ndarray[double, ndim=1] cdd = np.empty(n_vel_dims, dtype=np.float64)
 
     cdef np.ndarray[double, ndim=1] f = np.empty(n_vel_dims, dtype=np.float64)
+    cdef np.ndarray[double, ndim=1] goal_y_minus_start_y
 
     cdef int pps
     cdef int pvs
@@ -600,26 +601,36 @@ cpdef dmp_step_dual_cartesian(
             for ops, ovs in ((slice(3, 7), slice(3, 6)), (slice(10, 14), slice(9, 12))):
                 cdd[ovs] += p_gain * compact_axis_angle_from_quaternion(tracking_error[ops]) / dt
 
+        s = phase(current_t, forcing_term.alpha_z, goal_t, start_t, int_dt=int_dt)
+
         # position components
         for pps, pvs in POS_INDICES:
             current_ydd[pvs] = (
-                alpha_y * (beta_y * (goal_y[pps] - current_y[pps])
-                           + execution_time * (goal_yd[pvs] - current_yd[pvs]))
-                + f[pvs] + cdd[pvs]
-                ) / execution_time ** 2 + goal_ydd[pvs]
+                alpha_y * (
+                    beta_y * (goal_y[pps] - current_y[pps])
+                    + execution_time * (goal_yd[pvs] - current_yd[pvs])
+                    - beta_y * (goal_y[pps] - start_y[pps]) * s
+                )
+                + f[pvs]
+                + cdd[pvs]
+            ) / execution_time ** 2 + goal_ydd[pvs]
             current_yd[pvs] += dt * current_ydd[pvs] + cd[pvs] / execution_time
             current_y[pps] += dt * current_yd[pvs]
 
         # orientation components
         for ops, ovs in ((slice(3, 7), slice(3, 6)), (slice(10, 14), slice(9, 12))):
+            goal_y_minus_start_y = compact_axis_angle_from_quaternion(
+                concatenate_quaternions(goal_y[ops], q_conj(start_y[ops])))
             current_ydd[ovs] = (
                 alpha_y * (
-                    beta_y * compact_axis_angle_from_quaternion(
-                        concatenate_quaternions(
-                            goal_y[ops], q_conj(current_y[ops]))
-                    )
+                    beta_y * compact_axis_angle_from_quaternion(concatenate_quaternions(goal_y[ops], q_conj(current_y[ops])))
+                    + execution_time * goal_yd[ovs]
                     - execution_time * current_yd[ovs]
-                ) + f[ovs] + cdd[ovs]
+                    - beta_y * s * goal_y_minus_start_y
+                )
+                + goal_ydd[ovs] * execution_time ** 2
+                + f[ovs]
+                + cdd[ovs]
             ) / execution_time ** 2
             current_yd[ovs] += dt * current_ydd[ovs] + cd[ovs] / execution_time
             current_y[ops] = concatenate_quaternions(