diff --git a/python/iLQR_manipulator_CP.py b/python/iLQR_manipulator_CP.py
index bbd43d3c9a630ae6e547fd1157d9feb86686af89..463dc080e2e63af0e5068ff93757cc6e9a7266c2 100644
--- a/python/iLQR_manipulator_CP.py
+++ b/python/iLQR_manipulator_CP.py
@@ -31,64 +31,62 @@ from math import factorial
 # ===============================
 
 # Logarithmic map for R^2 x S^1 manifold
-def logmap(f,f0):
-	position_error = f[:,:2] - f0[:,:2]
-	orientation_error = np.imag(np.log( np.exp( f0[:,-1]*1j ).conj().T * np.exp(f[:,-1]*1j).T )).conj().reshape((-1,1))
-	error = np.hstack(( position_error , orientation_error ))
-	return error
-
-# Forward kinematics for E-E
-def fkin(param,x):
-	x = x.T
-	A = np.tril(np.ones([param.nbVarX,param.nbVarX]))
-	f = np.vstack((param.linkLengths @ np.cos(A @ x), 
-				   param.linkLengths @ np.sin(A @ x), 
-				   np.mod(np.sum(x,0)+np.pi, 2*np.pi) - np.pi)) #x1,x2,o (orientation as single Euler angle for planar robot)
-	return f.T
-
-# Forward Kinematics for all joints
-def fkin0(param,x):
-	T = np.tril(np.ones([param.nbVarX,param.nbVarX]))
-	T2 = np.tril(np.matlib.repmat(param.linkLengths,len(x),1))
-	f = np.vstack(( 
-		T2 @ np.cos(T@x),
-		T2 @ np.sin(T@x)
-	)).T
-	f = np.vstack(( 
-		np.zeros(2),
-		f
-	))
+def logmap(f, f0):
+	position_error = f[:2,:] - f0[:2,:]
+	orientation_error = np.imag(np.log(np.exp(f0[-1,:]*1j).conj().T * np.exp(f[-1,:]*1j).T)).conj()
+	diff = np.vstack([position_error, orientation_error])
+	return diff
+
+# Forward kinematics for end-effector (in robot coordinate system)
+def fkin(x, param):
+	L = np.tril(np.ones([param.nbVarX, param.nbVarX]))
+	f = np.vstack([
+		param.l @ np.cos(L @ x),
+		param.l @ np.sin(L @ x),
+		np.mod(np.sum(x,0)+np.pi, 2*np.pi) - np.pi
+	]) # f1,f2,f3, where f3 is the orientation (single Euler angle for planar robot)
+	return f
+
+# Forward kinematics for all joints (in robot coordinate system)
+def fkin0(x, param):
+	L = np.tril(np.ones([param.nbVarX, param.nbVarX]))
+	f = np.vstack([
+		L @ np.diag(param.l) @ np.cos(L @ x),
+		L @ np.diag(param.l) @ np.sin(L @ x)
+	])
+	f = np.hstack([np.zeros([2,1]), f])
 	return f
 
 # Jacobian with analytical computation (for single time step)
-def jkin(param,x):
-	T = np.tril( np.ones((len(x),len(x))) )
-	J = np.vstack((
-		-np.sin(T@x).T @ np.diag(param.linkLengths) @ T,
-		np.cos(T@x).T @ np.diag(param.linkLengths) @ T,
-		np.ones(len(x))
-	))
+def Jkin(x, param):
+	L = np.tril(np.ones([param.nbVarX, param.nbVarX]))
+	J = np.vstack([
+		-np.sin(L @ x).T @ np.diag(param.l) @ L,
+		 np.cos(L @ x).T @ np.diag(param.l) @ L,
+		 np.ones([1,param.nbVarX])
+	])
 	return J
 
-# Residual and Jacobian
-def f_reach(param,x):
-	f = logmap(fkin(param,x),param.mu)
-	J = np.zeros(( len(x) * param.nbVarF , len(x) * param.nbVarX ))
 
-	for t in range(x.shape[0]):
-		f[t,:2] = param.A[t].T @ f[t,:2] # Object oriented fk
-		Jtmp = jkin(param,x[t])
-		Jtmp[:2] = param.A[t].T @ Jtmp[:2] # Object centered jacobian
+# Residual and Jacobian for a viapoints reaching task (in object coordinate system)
+def f_reach(x, param):
+	f = logmap(fkin(x, param), param.Mu)
+	J = np.zeros([param.nbPoints * param.nbVarF, param.nbPoints * param.nbVarX])
+	for t in range(param.nbPoints):
+		f[:2, t] = param.A[:, :, t].T @ f[:2, t]  # Object oriented residual
+		Jtmp = Jkin(x[:, t], param)
+		Jtmp[:2] = param.A[:, :, t].T @ Jtmp[:2]  # Object centered Jacobian
 
 		if param.useBoundingBox:
 			for i in range(2):
-				if abs(f[t,i]) < param.sizeObj[i]:
-					f[t,i] = 0
-					Jtmp[i]=0
+				if abs(f[i, t]) < param.sz[i]:
+					f[i, t] = 0
+					Jtmp[i] = 0
 				else:
-					f[t,i] -=  np.sign(f[t,i]) * param.sizeObj[i]
-		J[ t*param.nbVarF:(t+1)*param.nbVarF , t*param.nbVarX:(t+1)*param.nbVarX] = Jtmp
-	return f,J
+					f[i, t] -= np.sign(f[i, t]) * param.sz[i]
+
+		J[t * param.nbVarF:(t + 1) * param.nbVarF, t * param.nbVarX:(t + 1) * param.nbVarX] = Jtmp
+	return f, J
 
 # Building piecewise constant basis functions
 def build_phi_piecewise(nb_data, nb_fct):
@@ -131,37 +129,31 @@ param.nbPoints = 2 # Number of viapoints
 param.nbVarX = 3 # State space dimension (x1,x2,x3)
 param.nbVarU = 3 # Control space dimension (dx1,dx2,dx3)
 param.nbVarF = 3 # Objective function dimension (f1,f2,f3, with f3 as orientation)
-param.linkLengths = [2,2,1] # Robot links lengths
-param.sizeObj = [.2,.3] # Size of objects
+param.l = [2,2,1] # Robot links lengths
+param.sz = [.2,.3] # Size of objects
 param.r = 1e-6 # Control weight term
-param.mu = np.asarray([[2,1,-np.pi/6], [3,2,-np.pi/3]]).T # Viapoints 
+param.Mu = np.asarray([[2,1,-np.pi/2], [3,1,-np.pi/2]]).T # Viapoints
 param.useBoundingBox = True # Consider bounding boxes for reaching cost
 param.nbFct = 5 # Number of basis functions
+param.A = np.zeros([2, 2, param.nbPoints]) # Object orientation matrices
 param.basisName = "PIECEWISE"
 
-# Task parameters
+# Main program
 # ===============================
 
-# Targets
-param.mu = np.asarray([
-	[2,1,-np.pi/2],				# x , y , orientation
-	[3,1,-np.pi/2]
-])
-
-# Transformation matrices
-param.A = np.zeros( (param.nbPoints,2,2) )
-for i in range(param.nbPoints):
-	orn_t = param.mu[i,-1]
-	param.A[i,:,:] = np.asarray([
-		[np.cos(orn_t) , -np.sin(orn_t)],
-		[np.sin(orn_t) , np.cos(orn_t)]
+# Object rotation matrices
+for t in range(param.nbPoints):
+	orn_t = param.Mu[-1,t]
+	param.A[:,:,t] = np.asarray([
+		[np.cos(orn_t), -np.sin(orn_t)],
+		[np.sin(orn_t), np.cos(orn_t)]
 	])
 
-# Regularization matrix
-R = np.identity( (param.nbData-1) * param.nbVarU ) * param.r
-
 # Precision matrix
-Q = np.identity( param.nbVarF  * param.nbPoints)
+Q = np.identity(param.nbVarF * param.nbPoints)
+
+# Control weight matrix
+R = np.identity((param.nbData-1) * param.nbVarU) * param.r
 
 # System parameters
 # ===============================
@@ -171,13 +163,13 @@ tl = np.linspace(0,param.nbData,param.nbPoints+1)
 tl = np.rint(tl[1:]).astype(np.int64)-1
 idx = np.array([ i + np.arange(0,param.nbVarX,1) for i in (tl* param.nbVarX)]) 
 
-u = np.zeros( param.nbVarU * (param.nbData-1) ) # Initial control command
-x0 = np.array( [3*np.pi/4 , -np.pi/2 , - np.pi/4] ) # Initial state (in joint space)
 
 # Transfer matrices (for linear system as single integrator)
-Su0 = np.vstack([np.zeros((param.nbVarX, param.nbVarX*(param.nbData-1))), 
-      np.tril(np.kron(np.ones((param.nbData-1, param.nbData-1)), np.eye(param.nbVarX)*param.dt))]) 
-Sx0 = np.kron( np.ones(param.nbData) , np.identity(param.nbVarX) ).T
+Su0 = np.vstack([
+	np.zeros([param.nbVarX, param.nbVarX*(param.nbData-1)]),
+	np.tril(np.kron(np.ones([param.nbData-1, param.nbData-1]), np.eye(param.nbVarX) * param.dt))
+])
+Sx0 = np.kron(np.ones(param.nbData), np.identity(param.nbVarX)).T
 Su = Su0[idx.flatten()] # We remove the lines that are out of interest
 
 # Basis functions
@@ -195,33 +187,32 @@ PSI = np.kron(phi,np.identity(param.nbVarU))
 # Solving iLQR
 # ===============================
 
-for i in range( param.nbIter ):
-	x = np.real(Su0 @ u + Sx0 @ x0)
-	x = x.reshape( (  param.nbData , param.nbVarX) )
+u = np.zeros( param.nbVarU * (param.nbData-1) ) # Initial control command
+x0 = np.array( [3*np.pi/4 , -np.pi/2 , - np.pi/4] ) # Initial state (in joint space)
 
-	f, J = f_reach(param,x[tl])
-	dw = np.linalg.inv(PSI.T @ Su.T @ J.T @ Q @ J @ Su @ PSI + PSI.T @ R @ PSI) @ (-PSI.T @ Su.T @ J.T @ Q @ f.flatten() - PSI.T @ u * param.r)
+for i in range( param.nbIter ):
+	x = Su0 @ u + Sx0 @ x0 # System evolution
+	x = x.reshape([param.nbVarX, param.nbData], order='F')
+	f, J = f_reach(x[:,tl], param) # Residuals and Jacobians
+	dw = np.linalg.inv(PSI.T @ Su.T @ J.T @ Q @ J @ Su @ PSI + PSI.T @ R @ PSI) @\
+		 (-PSI.T @ Su.T @ J.T @ Q @ f.flatten('F') - PSI.T @ u * param.r)
 	du = PSI @ dw
-	# Perform line search
+	# Estimate step size with backtracking line search method
 	alpha = 1
-	cost0 = f.flatten() @ Q @ f.flatten() + np.linalg.norm(u) * param.r
-	
+	cost0 = f.flatten('F').T @ Q @ f.flatten('F') + np.linalg.norm(u)**2 * param.r # Cost
 	while True:
 		utmp = u + du * alpha
-		xtmp = np.real(Su0 @ utmp + Sx0 @ x0)
-		xtmp = xtmp.reshape( (  param.nbData , param.nbVarX) )
-		ftmp, _ = f_reach(param,xtmp[tl])
-		cost = ftmp.flatten() @ Q @ ftmp.flatten() + np.linalg.norm(utmp) * param.r
-		
+		xtmp = Su0 @ utmp + Sx0 @ x0 # System evolution
+		xtmp = xtmp.reshape([param.nbVarX, param.nbData], order='F')
+		ftmp,_ = f_reach(xtmp[:,tl], param) # Residuals
+		cost = ftmp.flatten('F').T @ Q @ ftmp.flatten('F') + np.linalg.norm(utmp)**2 * param.r # Cost
 		if cost < cost0 or alpha < 1e-3:
 			u = utmp
-			print("Iteration {}, cost: {}, alpha: {}".format(i,cost,alpha))
+			print("Iteration {}, cost: {}".format(i,cost))
 			break
-
 		alpha /=2
-
-	if np.linalg.norm( alpha * du ) < 1e-2:
-		break
+	if np.linalg.norm(du * alpha) < 1E-2:
+		break # Stop iLQR iterations when solution is reached
 
 # Plotting
 # ===============================
@@ -231,17 +222,17 @@ plt.axis("off")
 plt.gca().set_aspect('equal', adjustable='box')
 
 # Get points of interest
-f = fkin(param,x)
-f00 = fkin0(param,x[0])
-f10 = fkin0(param,x[tl[0]])
-fT0 = fkin0(param,x[-1])
+f = fkin(x, param)
+f00 = fkin0(x[:, 0], param)
+f10 = fkin0(x[:, tl[0]], param)
+fT0 = fkin0(x[:, -1], param)
 u = u.reshape((-1,param.nbVarU))
 
-plt.plot( f00[:,0] , f00[:,1],c='black',linewidth=5,alpha=.2)
-plt.plot( f10[:,0] , f10[:,1],c='black',linewidth=5,alpha=.4)
-plt.plot( fT0[:,0] , fT0[:,1],c='black',linewidth=5,alpha=.6)
+plt.plot( f00[0, :] , f00[1,:],c='black',linewidth=5,alpha=.2)
+plt.plot( f10[0, :] , f10[1,:],c='black',linewidth=5,alpha=.4)
+plt.plot( fT0[0, :] , fT0[1,:],c='black',linewidth=5,alpha=.6)
 
-plt.plot(f[:,0],f[:,1],c="black",marker="o",markevery=[0]+tl.tolist()) #,label="Trajectory"2
+plt.plot(f[0, :],f[1, :],c="black",marker="o",markevery=[0]+tl.tolist()) #,label="Trajectory"2
 
 # Plot bounding box or via-points
 ax = plt.gca()
@@ -249,36 +240,36 @@ color_map = ["deepskyblue","darkorange"]
 for i in range(param.nbPoints):
 	
 	if param.useBoundingBox:
-		rect_origin = param.mu[i,:2] - param.A[i]@np.array(param.sizeObj)
-		rect_orn = param.mu[i,-1]
+		rect_origin = param.Mu[i,:2] - param.A[i]@np.array(param.sz)
+		rect_orn = param.Mu[i,-1]
 
-		rect = patches.Rectangle(rect_origin,param.sizeObj[0]*2,param.sizeObj[1]*2,np.degrees(rect_orn),color=color_map[i])
+		rect = patches.Rectangle(rect_origin,param.sz[0]*2,param.sz[1]*2,np.degrees(rect_orn),color=color_map[i])
 		ax.add_patch(rect)
 	else:
-		plt.scatter(param.mu[i,0],param.mu[i,1],s=100,marker="X",c=color_map[i])
+		plt.scatter(param.Mu[i,0],param.Mu[i,1],s=100,marker="X",c=color_map[i])
 
 fig,axs = plt.subplots(7,1)
 
-axs[0].plot(f[:,0],c='black')
+axs[0].plot(f[0],c='black')
 axs[0].set_ylabel("$f(x)_1$")
 axs[0].set_xticks([0,param.nbData])
 axs[0].set_xticklabels(["0","T"])
 for i in range(param.nbPoints):
-    axs[0].scatter(tl[i],param.mu[i,0],c='blue')
+    axs[0].scatter(tl[i],param.Mu[0, i],c='blue')
 
-axs[1].plot(f[:,1],c='black')
+axs[1].plot(f[1],c='black')
 axs[1].set_ylabel("$f(x)_2$")
 axs[1].set_xticks([0,param.nbData])
 axs[1].set_xticklabels(["0","T"])
 for i in range(param.nbPoints):
-    axs[1].scatter(tl[i],param.mu[i,1],c='blue')
+    axs[1].scatter(tl[i],param.Mu[1, i],c='blue')
 
-axs[2].plot(f[:,2],c='black')
+axs[2].plot(f[2],c='black')
 axs[2].set_ylabel("$f(x)_3$")
 axs[2].set_xticks([0,param.nbData])
 axs[2].set_xticklabels(["0","T"])
 for i in range(param.nbPoints):
-    axs[2].scatter(tl[i],param.mu[i,2],c='blue')
+    axs[2].scatter(tl[i],param.Mu[2, i],c='blue')
 
 axs[3].plot(u[:,0],c='black')
 axs[3].set_ylabel("$u_1$")