def callback(self, msg):
if self.js_prev is not None:
self.js_prev = msg
self.js_current = msg
else:
self.js_prev = self.js_current
self.js_current = msg
# If your joints have changed by more than a small threshold amount, then you are moving! Start recording joint angles
if np.sum(
np.abs(
np.array(self.js_prev.position) -
np.array(self.js_current.position))) > 0.01:
print('Moving')
j_list = [
msg.position[msg.name.index(j)] for j in self.joint_names
]
self.traj.append(j_list)
elif len(self.traj) > 0:
# Fit DMP to motion segment
path = np.array(self.traj)
dmp = DMP(path[0, :],
path[-1, :],
num_basis_funcs=500,
dt=0.01,
d=path.shape[1],
jnames=self.joint_names)
_, weights = imitate_path(
path + 1e-5 * np.random.randn(path.shape[0], path.shape[1]),
dmp)
print("Learned weights: {}".format(weights))
dmp.weights = weights
self.dmp_count += 1
# Save it with pickle
with open("./dmp{}.npy".format(self.dmp_count), "w") as f:
pickle.dump(dmp, f, pickle.HIGHEST_PROTOCOL)
# Plot the path imitated by the learned dmp (rolled out) against the actual trajectory path...
y_r, dy_r, _ = dmp.rollout()
plot_rollout(y_r, path)
jt = build_jt(msg.header, self.joint_names, y_r, dy_r)
self.traj_pub.publish(jt)
self.traj = []
else:
print("Stationary")
