Add batch option to iLQR_manipulator_obstacle.m

matlab/iLQR_manipulator_obstacle.m

@@ -2,7 +2,8 @@
 %% 	  (viapoints task with position+orientation including bounding boxes on f(x))
 %%
 %%    Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
-%%    Written by Sylvain Calinon <https://calinon.ch>
+%%    Written by Teguh Lembono <teguh.lembono@idiap.ch>, 
+%%		Sylvain Calinon <https://calinon.ch>
 %%
 %%    This file is part of RCFS.
 %%
@@ -24,7 +25,7 @@ function iLQR_manipulator_obstacle
 %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
 param.dt = 1E-2; %Time step size
 param.nbData = 50; %Number of datapoints
-param.nbIter = 50; %Number of iterations for iLQR
+param.nbIter = 30; %Number of iterations for iLQR
 param.nbPoints = 1; %Number of viapoints
 param.nbObstacles = 2; %Number of obstacles
 param.nbVarX = 3; %State space dimension (q1,q2,q3)
@@ -35,6 +36,7 @@ param.sz2 = [.5, .8]; %Size of obstacles
 param.q = 1E2; %Tracking weight term
 param.q2 = 1E0; %Obstacle avoidance weight term
 param.r = 1E-3; %Control weight term
+param.useBatch = true; %use batch computation for collision cost
 
 param.Mu = [[3; -1; 0]]; %Viapoints (x1,x2,o)
 
@@ -64,11 +66,16 @@ Su0 = [zeros(param.nbVarX, param.nbVarX*(param.nbData-1)); tril(kron(ones(param.
 Sx0 = kron(ones(param.nbData,1), eye(param.nbVarX));
 Su = Su0(idx,:);
 
+tic;
 for n=1:param.nbIter
 	x = reshape(Su0 * u + Sx0 * x0, param.nbVarX, param.nbData); %System evolution
 	[f, J] = f_reach(x(:,tl), param); %Tracking objective
 	
-	[f2, J2, id2] = f_avoid(x, param); %Avoidance objective
+	if param.useBatch
+		[f2, J2, id2] = f_avoid2(x, param); %Avoidance objective
+	else
+		[f2, J2, id2] = f_avoid(x, param); %Avoidance objective
+	end
 	Su2 = Su0(id2,:);
 	
 	du = (Su' * J' * J * Su * param.q + Su2' * J2' * J2 * Su2 * param.q2 + R) \ ...
@@ -81,7 +88,11 @@ for n=1:param.nbIter
 		utmp = u + du * alpha;
 		xtmp = reshape(Su0 * utmp + Sx0 * x0, param.nbVarX, param.nbData);
 		ftmp = f_reach(xtmp(:,tl), param); %Tracking objective
-		ftmp2 = f_avoid(xtmp, param); %Avoidance objective
+		if param.useBatch
+			ftmp2 = f_avoid2(xtmp, param); %Avoidance objective
+		else
+			ftmp2 = f_avoid(xtmp, param); %Avoidance objective
+		end
 		cost = norm(ftmp(:))^2 * param.q + norm(ftmp2(:))^2 * param.q2 + norm(utmp)^2 * param.r;
 		if cost < cost0 || alpha < 1E-3
 			break;
@@ -89,13 +100,14 @@ for n=1:param.nbIter
 		alpha = alpha * 0.5;
 	end
 	u = u + du * alpha;
-	
+	disp(cost)
 	if norm(du * alpha) < 1E-2
 		break; %Stop iLQR when solution is reached
 	end
 end
 disp(['iLQR converged in ' num2str(n) ' iterations.']);
-
+toc
+%disp(tEnd-tStart)
 %Log data
 r.x = x;
 r.f = [];
@@ -192,6 +204,7 @@ function [f, J, id] = f_avoid(x, param)
 					f = [f; ftmp];
 					Jrob = [jkin(x(1:j,t), param.l(1:j)), zeros(param.nbVarF, param.nbVarU-j)];
 					Jtmp = -e' * param.U2(:,:,i)' * Jrob(1:2,:); %quadratic form
+
 					J = blkdiag(J, Jtmp);
 					id = [id, (t-1) * param.nbVarU + [1:param.nbVarU]'];
 				end
@@ -199,3 +212,46 @@ function [f, J, id] = f_avoid(x, param)
 		end
 	end
 end
+
+%%%%%%%%%%%%%%%%%%%%%%
+%Forward kinematics for all robot articulations in batch
+function f = fkin0batch(x, L)
+	T = tril(ones(size(x,1)));
+	T2 = T * diag(L);
+	f = [(T2 * cos(T * x))(:), ...
+			(T2 * sin(T * x))(:)]; 
+end
+
+%%%%%%%%%%%%%%%%%%%%%%
+%Cost and gradient for an obstacle avoidance task with batch computation
+function [f, J, id] = f_avoid2(x, param)
+	f=[]; J=[]; id=[];
+	for i=1:param.nbObstacles
+		xee_batch = fkin0batch(x, param.l);
+		n= size(xee_batch,1);
+		e = (xee_batch - repmat(param.Mu2(1:2,i)', n , 1)) * param.U2(:,:,i);
+		ftmp = ones(n,1) - sum(e.^2, 2); %quadratic form
+		
+		%get the indices where there is collision (ftmp > 0)
+		idx_active = linspace(1,n,n)(ftmp>0);
+		
+		for idx=1:size(idx_active,2)
+			j = idx_active(idx);
+			t = floor((j-1)/param.nbVarX)+1;
+			k = mod(j-1,param.nbVarX)+1;	
+	
+			%compute Jacobian
+			Jrob = [jkin(x(1:k,t), param.l(1:k)), zeros(param.nbVarF, param.nbVarX-k)];
+			Jtmp = -e(j,:) * param.U2(:,:,i)' * Jrob(1:2,:); %quadratic form
+			Jj = zeros(1,param.nbVarX*param.nbData);
+			Jj(:,param.nbVarX*(t-1)+1:param.nbVarX*(t)) = Jtmp;
+			J = [J; Jj];
+			id = [id, (t-1) * param.nbVarX + [1:param.nbVarX]'];
+		end
+		
+		f = [f; ftmp(idx_active(:))];		#keep only active value
+	end
+	
+	id = unique(id); #remove duplicate
+	J = J(:,id); #keep only active value
+end
\ No newline at end of file

python/iLQR_manipulator_obstacle.py

@@ -87,7 +87,7 @@ def f_avoid(x, param):
 
                 if ftmp > 0:
                     f += [ftmp]
-                    J_rbt = np.hstack(( jkin(x[t,:(j+1)] , param.linkLengths[:(j+1)]), np.zeros(( param.nbVarF , param.nbVarU-(j+1))) ))
+                    J_rbt = np.hstack(( jkin(x[t,:(j+1)] , param.linkLengths[:(j+1)]), np.zeros(( param.nbVarF , param.nbVarX-(j+1))) ))
                     J_tmp = (-e.T @ param.U_obs[i].T @ J_rbt[:2]).reshape((-1,1))
                     #J_tmp = 1*(J_rbt[:2].T @ param.U_obs[i] @ e).reshape((-1,1))
 
@@ -96,7 +96,7 @@ def f_avoid(x, param):
                     J2[-J_tmp.shape[0]:,-J_tmp.shape[1]:] = J_tmp
                     J = J2 # Numpy does not provid a blockdiag function...
 
-                    idx.append( t*param.nbVarU + np.array(range(param.nbVarU)) )
+                    idx.append( t*param.nbVarX + np.array(range(param.nbVarX)) )
 
     f = np.array(f)
     idx = np.array(idx)