Created optistack rotation trajectory generation to use fatrop and fi…

…le to compare with other OCP specifications;fatrop not working yet

examples/trajectory_generation/COMPARING trajectory_generation_fullpose.py

@@ -0,0 +1,262 @@
+
+# Imports
+import numpy as np
+from math import pi
+import invariants_py.data_handler as dh
+import matplotlib.pyplot as plt
+import invariants_py.reparameterization as reparam
+import scipy.interpolate as ip
+from invariants_py.calculate_invariants.opti_calculate_vector_invariants_rotation import OCP_calc_rot
+from invariants_py.calculate_invariants.opti_calculate_vector_invariants_position import OCP_calc_pos as OCP_calc_pos
+from invariants_py.generate_trajectory.opti_generate_orientation_traj_from_vector_invars import OCP_gen_rot
+from invariants_py.generate_trajectory.opti_generate_position_traj_from_vector_invars import OCP_gen_pos as OCP_gen_pos
+from invariants_py.generate_trajectory.opti_generate_pose_traj_from_vector_invars_fatrop import OCP_gen_pose as OCP_gen_pose_fatrop
+from invariants_py.generate_trajectory.rockit_generate_pose_traj_from_vector_invars import OCP_gen_pose as OCP_gen_pose_rockit
+from IPython.display import clear_output
+from scipy.spatial.transform import Rotation as R
+from invariants_py.kinematics.rigidbody_kinematics import orthonormalize_rotation as orthonormalize
+from stl import mesh
+import invariants_py.plotting_functions.plotters as pl
+from invariants_py.ocp_initialization import initial_trajectory_movingframe_rotation
+#%%
+data_location = dh.find_data_path('beer_1.txt')
+opener_location =  dh.find_data_path('opener.stl')
+bottle_location = dh.find_data_path('bottle.stl')
+
+trajectory,time = dh.read_pose_trajectory_from_data(data_location,dtype = 'txt')
+pose,time_profile,arclength,nb_samples,stepsize = reparam.reparameterize_trajectory_arclength(trajectory)
+arclength_n = arclength/arclength[-1]
+home_pos = [0,0,0] # Use this if not considering the robot
+# home_pos = [0.3056, 0.0635, 0.441] # Define home position of the robot
+trajectory_position = pose[:,:3,3] + home_pos
+trajectory_orientation = pose[:,:3,:3]
+
+fig = plt.figure(figsize=(8,8))
+ax = fig.add_subplot(111, projection='3d')
+ax = plt.axes(projection='3d')
+ax.plot(trajectory_position[:,0],trajectory_position[:,1],trajectory_position[:,2],'.-')
+n_frames = 10
+indx = np.trunc(np.linspace(0,len(trajectory_orientation)-1,n_frames))
+indx = indx.astype(int)
+
+for i in indx:
+    pl.plot_3d_frame(trajectory_position[i,:],trajectory_orientation[i,:,:],1,0.01,['red','green','blue'],ax)
+    pl.plot_stl(opener_location,trajectory_position[i,:],trajectory_orientation[i,:,:],colour="c",alpha=0.2,ax=ax)    
+
+#%%
+# define class for OCP results
+class OCP_results:
+
+    def __init__(self,FSt_frames,FSr_frames,Obj_pos,Obj_frames,invariants):
+        self.FSt_frames = FSt_frames
+        self.FSr_frames = FSr_frames
+        self.Obj_pos = Obj_pos
+        self.Obj_frames = Obj_frames
+        self.invariants = invariants
+
+optim_calc_results = OCP_results(FSt_frames = [], FSr_frames = [], Obj_pos = [], Obj_frames = [], invariants = np.zeros((len(trajectory),6)))
+
+# specify optimization problem symbolically
+FS_calculation_problem_pos = OCP_calc_pos(window_len=nb_samples, bool_unsigned_invariants = False, rms_error_traj = 0.004)
+FS_calculation_problem_rot = OCP_calc_rot(window_len=nb_samples, bool_unsigned_invariants = False, rms_error_traj = 2*pi/180) 
+
+# calculate invariants given measurements
+optim_calc_results.invariants[:,3:], optim_calc_results.Obj_pos, optim_calc_results.FSt_frames = FS_calculation_problem_pos.calculate_invariants(trajectory,stepsize)
+optim_calc_results.invariants[:,:3], optim_calc_results.Obj_frames, optim_calc_results.FSr_frames = FS_calculation_problem_rot.calculate_invariants(trajectory,stepsize)
+optim_calc_results.Obj_pos += home_pos
+
+fig = plt.figure(figsize=(8,8))
+ax = fig.add_subplot(111, projection='3d')
+ax = plt.axes(projection='3d')
+ax.plot(trajectory_position[:,0],trajectory_position[:,1],trajectory_position[:,2],'b')
+ax.plot(optim_calc_results.Obj_pos[:,0],optim_calc_results.Obj_pos[:,1],optim_calc_results.Obj_pos[:,2],'r')
+indx = np.trunc(np.linspace(0,len(optim_calc_results.Obj_pos)-1,n_frames))
+indx = indx.astype(int)
+for i in indx:
+    pl.plot_3d_frame(optim_calc_results.Obj_pos[i,:],optim_calc_results.Obj_frames[i,:,:],1,0.01,['red','green','blue'],ax)
+    pl.plot_stl(opener_location,trajectory_position[i,:],trajectory_orientation[i,:,:],colour="c",alpha=0.2,ax=ax)
+    pl.plot_stl(opener_location,optim_calc_results.Obj_pos[i,:],optim_calc_results.Obj_frames[i,:,:],colour="r",alpha=0.2,ax=ax)
+
+pl.plot_orientation(optim_calc_results.Obj_frames,trajectory_orientation)
+
+pl.plot_invariants(optim_calc_results.invariants,[],arclength_n)
+
+if plt.get_backend() != 'agg':
+    plt.show()
+
+#%%
+# Spline of model
+knots = np.concatenate(([arclength_n[0]],[arclength_n[0]],arclength_n,[arclength_n[-1]],[arclength_n[-1]]))
+degree = 3
+spline_model_trajectory = ip.BSpline(knots,optim_calc_results.invariants,degree)
+
+def interpolate_model_invariants(demo_invariants, progress_values):
+
+    resampled_invariants = np.array([demo_invariants(i) for i in progress_values]) 
+    new_stepsize = progress_values[1] - progress_values[0] 
+
+    resampled_invariants[:,0] = resampled_invariants[:,0] *  (progress_values[-1] - progress_values[0])
+    return resampled_invariants, new_stepsize
+
+#%% 
+current_progress = 0
+number_samples = 100
+
+progress_values = np.linspace(current_progress, arclength_n[-1], number_samples)
+model_invariants,new_stepsize = interpolate_model_invariants(spline_model_trajectory,progress_values)
+
+pl.plot_interpolated_invariants(optim_calc_results.invariants, model_invariants, arclength_n, progress_values)
+
+# new constraints
+current_index = round(current_progress*len(trajectory))
+p_obj_start = optim_calc_results.Obj_pos[current_index]
+R_obj_start = orthonormalize(optim_calc_results.Obj_frames[current_index])
+FSt_start = orthonormalize(optim_calc_results.FSt_frames[current_index])
+FSr_start = orthonormalize(optim_calc_results.FSr_frames[current_index])
+p_obj_end = optim_calc_results.Obj_pos[-1] + np.array([0.1,0.1,0.1])
+rotate = R.from_euler('z', 0, degrees=True)
+R_obj_end =  orthonormalize(rotate.apply(optim_calc_results.Obj_frames[-1]))
+FSt_end = orthonormalize(optim_calc_results.FSt_frames[-1])
+FSr_end = orthonormalize(optim_calc_results.FSr_frames[-1])
+
+# define new class for OCP results
+optim_gen_results = OCP_results(FSt_frames = [], FSr_frames = [], Obj_pos = [], Obj_frames = [], invariants = np.zeros((number_samples,6)))
+fatrop_gen_results = OCP_results(FSt_frames = [], FSr_frames = [], Obj_pos = [], Obj_frames = [], invariants = np.zeros((number_samples,6)))
+rockit_gen_results = OCP_results(FSt_frames = [], FSr_frames = [], Obj_pos = [], Obj_frames = [], invariants = np.zeros((number_samples,6)))
+
+# Linear initialization
+R_obj_init = reparam.interpR(np.linspace(0, 1, len(optim_calc_results.Obj_frames)), [0,1], np.array([R_obj_start, R_obj_end]))
+
+R_r_init, R_r_init_array, invars_init = initial_trajectory_movingframe_rotation(R_obj_start, R_obj_end)
+
+boundary_constraints = {
+    "position": {
+        "initial": p_obj_start,
+        "final": p_obj_end
+    },
+    "orientation": {
+        "initial": R_obj_start,
+        "final": R_obj_end
+    },
+    "moving-frame": {
+        "translational": {
+            "initial": FSt_start,
+            "final": FSt_end
+        },
+        "rotational": {
+            "initial": R_r_init,
+            "final": R_r_init
+        }
+    },
+}
+
+# Define OCP weights
+weights_params = {
+    "w_invars": np.array([1, 1, 1, 5*10**1, 1.0, 1.0]),
+    "w_high_start": 60,
+    "w_high_end": number_samples,
+    "w_high_invars": 10*np.array([1, 1, 1, 5*10**1, 1.0, 1.0]),
+    "w_high_active": 0
+}
+
+# specify optimization problem symbolically
+FS_online_generation_problem_pos = OCP_gen_pos(N=number_samples,w_invars = weights_params['w_invars'][3:])
+FS_online_generation_problem_rot = OCP_gen_rot(window_len=number_samples,w_invars = weights_params['w_invars'][:3])
+fatrop_online_generation_problem = OCP_gen_pose_fatrop(boundary_constraints, number_samples, solver = 'ipopt')
+rockit_online_generation_problem = OCP_gen_pose_rockit(boundary_constraints, number_samples,False)
+
+initial_values = {
+    "trajectory": {
+        "position": optim_calc_results.Obj_pos,
+        "orientation": R_obj_init
+    },
+    "moving-frame": {
+        "translational": optim_calc_results.FSt_frames,
+        "rotational": R_r_init_array,
+    },
+    "invariants": model_invariants,
+    # "joint-values": robot_params["q_init"] if robot_params["urdf_file_name"] is not None else {}
+}
+
+# Solve
+optim_gen_results.invariants[:,3:], optim_gen_results.Obj_pos, optim_gen_results.FSt_frames = FS_online_generation_problem_pos.generate_trajectory(U_demo = model_invariants[:,3:], p_obj_init = optim_calc_results.Obj_pos, R_t_init = optim_calc_results.FSt_frames, R_t_start = FSt_start, R_t_end = FSt_end, p_obj_start = p_obj_start, p_obj_end = p_obj_end, step_size = new_stepsize)
+optim_gen_results.invariants[:,:3], optim_gen_results.Obj_frames, optim_gen_results.FSr_frames = FS_online_generation_problem_rot.generate_trajectory(U_demo = model_invariants[:,:3], R_obj_init = optim_calc_results.Obj_frames, R_r_init = optim_calc_results.FSr_frames, R_r_start = FSr_start, R_r_end = FSr_end, R_obj_start = R_obj_start, R_obj_end = R_obj_end, step_size = new_stepsize)
+fatrop_gen_results.invariants, fatrop_gen_results.Obj_pos, fatrop_gen_results.Obj_frames, fatrop_gen_results.FSt_frames, fatrop_gen_results.FSr_frames = fatrop_online_generation_problem.generate_trajectory(model_invariants,boundary_constraints,new_stepsize,weights_params,initial_values)
+rockit_gen_results.invariants, rockit_gen_results.Obj_pos, rockit_gen_results.Obj_frames, rockit_gen_results.FSt_frames, rockit_gen_results.FSr_frames, tot_time, joint_values = rockit_online_generation_problem.generate_trajectory(model_invariants,boundary_constraints,new_stepsize,weights_params,initial_values)
+
+# for i in range(len(optim_gen_results.Obj_frames)):
+#     optim_gen_results.Obj_frames[i] = orthonormalize(optim_gen_results.Obj_frames[i])
+
+fig = plt.figure(figsize=(14,8))
+ax = fig.add_subplot(111, projection='3d')
+ax.plot(optim_calc_results.Obj_pos[:,0],optim_calc_results.Obj_pos[:,1],optim_calc_results.Obj_pos[:,2],'b')
+ax.plot(optim_gen_results.Obj_pos[:,0],optim_gen_results.Obj_pos[:,1],optim_gen_results.Obj_pos[:,2],'r')
+ax.plot(fatrop_gen_results.Obj_pos[:,0],fatrop_gen_results.Obj_pos[:,1],fatrop_gen_results.Obj_pos[:,2],'g')
+ax.plot(rockit_gen_results.Obj_pos[:,0],rockit_gen_results.Obj_pos[:,1],rockit_gen_results.Obj_pos[:,2],'m')
+for i in indx:
+    pl.plot_stl(opener_location,optim_calc_results.Obj_pos[i,:],optim_calc_results.Obj_frames[i,:,:],colour="c",alpha=0.2,ax=ax)
+
+indx_online = np.trunc(np.linspace(0,len(optim_gen_results.Obj_pos)-1,n_frames))
+indx_online = indx_online.astype(int)
+for i in indx_online:
+    # pl.plot_3d_frame(optim_calc_results.Obj_pos[i,:],optim_calc_results.Obj_frames[i,:,:],1,0.01,['red','green','blue'],ax)
+    # pl.plot_3d_frame(optim_gen_results.Obj_pos[i,:],optim_gen_results.Obj_frames[i,:,:],1,0.01,['red','green','blue'],ax)
+    pl.plot_stl(opener_location,optim_gen_results.Obj_pos[i,:],optim_gen_results.Obj_frames[i,:,:],colour="r",alpha=0.2,ax=ax)
+    pl.plot_stl(opener_location,fatrop_gen_results.Obj_pos[i,:],fatrop_gen_results.Obj_frames[i,:,:],colour="g",alpha=0.2,ax=ax)
+    pl.plot_stl(opener_location,rockit_gen_results.Obj_pos[i,:],rockit_gen_results.Obj_frames[i,:,:],colour="m",alpha=0.2,ax=ax)
+pl.plot_orientation(optim_calc_results.Obj_frames,optim_gen_results.Obj_frames,current_index)
+pl.plot_orientation(optim_calc_results.Obj_frames,fatrop_gen_results.Obj_frames,current_index)
+pl.plot_orientation(optim_calc_results.Obj_frames,rockit_gen_results.Obj_frames,current_index)
+
+
+fig = plt.figure()
+plt.subplot(2,3,1)
+plt.plot(arclength_n,optim_calc_results.invariants[:,0],'b')
+plt.plot(progress_values,optim_gen_results.invariants[:,0],'r')
+plt.plot(progress_values,fatrop_gen_results.invariants[:,0],'g')
+plt.plot(progress_values,rockit_gen_results.invariants[:,0],'m')
+plt.plot(0,0)
+plt.title('i_r1')
+
+plt.subplot(2,3,2)
+plt.plot(arclength_n,optim_calc_results.invariants[:,1],'b')
+plt.plot(progress_values,(optim_gen_results.invariants[:,1]),'r')
+plt.plot(progress_values,(fatrop_gen_results.invariants[:,1]),'g')
+plt.plot(progress_values,(rockit_gen_results.invariants[:,1]),'m')
+plt.plot(0,0)
+plt.title('i_r2')
+
+plt.subplot(2,3,3)
+plt.plot(arclength_n,optim_calc_results.invariants[:,2],'b')
+plt.plot(progress_values,(optim_gen_results.invariants[:,2]),'r')
+plt.plot(progress_values,(fatrop_gen_results.invariants[:,2]),'g')
+plt.plot(progress_values,(rockit_gen_results.invariants[:,2]),'m')
+plt.plot(0,0)
+plt.title('i_r3')
+
+plt.subplot(2,3,4)
+plt.plot(arclength_n,optim_calc_results.invariants[:,0],'b')
+plt.plot(progress_values,optim_gen_results.invariants[:,0],'r')
+plt.plot(arclength_n,fatrop_gen_results.invariants[:,0],'g')
+plt.plot(arclength_n,rockit_gen_results.invariants[:,0],'m')
+plt.plot(0,0)
+plt.title('i_t1')
+
+plt.subplot(2,3,5)
+plt.plot(arclength_n,optim_calc_results.invariants[:,1],'b')
+plt.plot(progress_values,(optim_gen_results.invariants[:,1]),'r')
+plt.plot(arclength_n,fatrop_gen_results.invariants[:,1],'g')
+plt.plot(arclength_n,rockit_gen_results.invariants[:,1],'m')
+plt.plot(0,0)
+plt.title('i_t1')
+
+plt.subplot(2,3,6)
+plt.plot(arclength_n,optim_calc_results.invariants[:,2],'b')
+plt.plot(progress_values,(optim_gen_results.invariants[:,2]),'r')
+plt.plot(arclength_n,fatrop_gen_results.invariants[:,2],'g')
+plt.plot(arclength_n,rockit_gen_results.invariants[:,2],'m')
+plt.plot(0,0)
+plt.title('i_t3')
+
+if plt.get_backend() != 'agg':
+    plt.show()

invariants_py/dynamics_vector_invariants.py

@@ -256,4 +256,25 @@ def define_integrator_invariants_rotation(h):
     out_plus1 = cas.vertcat(cas.vec(R_r_plus1),  cas.vec(R_obj_plus1))
     integrator = cas.Function("phi", [x,u,h] , [out_plus1])
 
+    return integrator
+
+def define_integrator_invariants_pose(h):
+    """Define a CasADi function that integrates the vector invariants for pose over a time interval h."""
+
+    # System states
+    R_t  = cas.MX.sym('R_t',3,3) # translational Frenet-Serret frame
+    p_obj = cas.MX.sym('p_obj',3,1) # object position
+    R_r  = cas.MX.sym('R_r',3,3) # rotational Frenet-Serret frame
+    R_obj = cas.MX.sym('R_obj',3,3) # object orientation
+    x = cas.vertcat(cas.vec(R_r), cas.vec(R_obj),cas.vec(R_t), p_obj)
+
+    # System controls (invariants)
+    u = cas.MX.sym('i',6,1)
+
+    ## Integrate symbolically and create a casadi function with inputs/outputs
+    (R_t_plus1, p_obj_plus1) = integrate_vector_invariants_position(R_t, p_obj, u[3:], h)
+    (R_r_plus1, R_obj_plus1) = integrate_vector_invariants_rotation(R_r, R_obj, u[:3], h)
+    out_plus1 = cas.vertcat(cas.vec(R_r_plus1),  cas.vec(R_obj_plus1),cas.vec(R_t_plus1),  p_obj_plus1)
+    integrator = cas.Function("phi", [x,u,h] , [out_plus1])
+
     return integrator