Add batch version for iLQR_manipulator_obstacle.py

python/iLQR_manipulator_obstacle.py

@@ -3,6 +3,7 @@
 
     Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
     Written by Jeremy Maceiras <jeremy.maceiras@idiap.ch>,
+    Teguh Lembono <teguh.lembono@idiap.ch>,
     Sylvain Calinon <https://calinon.ch>
 
     This file is part of RCFS.
@@ -24,6 +25,8 @@ import numpy as np
 import numpy.matlib
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
+import time
+
 
 # Helper functions
 # ===============================
@@ -37,12 +40,12 @@ def logmap(f,f0):
 
 # Forward kinematics for E-E
 def fkin(x,ls):
-	x = x.T
-	A = np.tril(np.ones([len(ls),len(ls)]))
-	f = np.vstack((ls @ np.cos(A @ x), 
-				   ls @ 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
+    x = x.T
+    A = np.tril(np.ones([len(ls),len(ls)]))
+    f = np.vstack((ls @ np.cos(A @ x), 
+                    ls @ 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(x):
@@ -58,6 +61,18 @@ def fkin0(x):
         ))
     return f
 
+# Forward Kinematics for all joints in batch
+def fkin0batch(x):
+    global param
+    T = np.tril(np.ones([param.nbVarX,param.nbVarX]))
+    T2 = np.tril(np.matlib.repmat(param.linkLengths,x.shape[1],1))
+    f = np.vstack(( 
+        (T2 @ np.cos(T@x.T)).flatten(order='F'),
+        (T2 @ np.sin(T@x.T)).flatten(order='F')
+        )).T
+    return f
+
+# Cost and jacobians for avoiding obstacles
 def f_avoid(x):
     global param
     f=[]
@@ -83,12 +98,47 @@ def f_avoid(x):
                     J2[-J_tmp.shape[0]:,-J_tmp.shape[1]:] = J_tmp
                     J = J2 # Numpy does not provid a blockdiag function...
 
-                    idx.append( (t-1)*param.nbVarU + np.array(range(param.nbVarU)) )
+                    idx.append( t*param.nbVarU + np.array(range(param.nbVarU)) )
 
     f = np.array(f)
     idx = np.array(idx)
     return f,J.T,idx
 
+# Cost and jacobians for avoiding obstacles, computed in batch
+def f_avoid2(x):
+    global param
+    f, J, idx = [], [], []
+    
+    for i in range(param.nbObstacles):
+        f_x = fkin0batch(x)                
+        e = (f_x- param.Obs[i][:2])@param.U_obs[i]        
+        ftmp = 1 - np.linalg.norm(e, axis=1)**2
+        
+        #check for value larger than zeros
+        idx_active = np.arange(len(ftmp))[ftmp>0]
+        
+        #compute Jacobian for each idx_active
+        for j in idx_active:
+            t = j//param.nbVarX # which time step
+            k = j%param.nbVarX  # which link
+            
+            J_rbt = np.hstack((jkin(x[t,:(k+1)] , param.linkLengths[:(k+1)]), np.zeros(( param.nbVarF , param.nbVarX-(k+1))) ))
+            J_tmp = (-e[j].T @ param.U_obs[i].T @ J_rbt[:2])
+            J_j = np.zeros(param.nbVarX*param.nbData)
+            J_j[param.nbVarX*t: param.nbVarX*(t+1)] = J_tmp
+            J += [J_j]            
+            
+            #Get indices
+            idx += np.arange(param.nbVarX*t, param.nbVarX*(t+1)).tolist()
+        f += [ftmp[idx_active]]
+    
+    idx = np.array(list(set(idx)))
+    f = np.concatenate(f)
+    if len(idx) > 0:
+        J = np.vstack(J)
+        J = J[:,idx]
+    return f, J,idx
+
 def jkin(x,ls):
     global param
     T = np.tril( np.ones((len(ls),len(ls))) )
@@ -127,6 +177,7 @@ param.sizeObstacle = [.5,.8] # Size of objects
 param.r = 1e-3 # Control weight term
 param.Q_track = 1e2
 param.Q_avoid = 1e0
+param.useBatch = True #Batch computation for collision avoidance cost
 
 # Task parameters
 # ===============================
@@ -183,13 +234,16 @@ Su = Su0[idx.flatten()] # We remove the lines that are out of interest
 
 # Solving iLQR
 # ===============================
-
+tic=time.time()
 for i in range( param.nbIter ):
     x = Su0 @ u + Sx0 @ x0
     x = x.reshape( (  param.nbData , param.nbVarX) )
 
     f, J = f_reach(x[tl])  # Tracking objective
-    f2, J2, id2 = f_avoid(x)# Avoidance objective
+    if param.useBatch:
+        f2, J2, id2 = f_avoid2(x)# Avoidance objective
+    else:
+        f2, J2, id2 = f_avoid(x)# Avoidance objective
 
     if len(id2) > 0: # Numpy does not allow zero sized array as Indices
         Su2 = Su0[id2.flatten()]
@@ -207,7 +261,11 @@ for i in range( param.nbIter ):
         xtmp = Su0 @ utmp + Sx0 @ x0
         xtmp = xtmp.reshape( (  param.nbData , param.nbVarX) )
         ftmp, _ = f_reach(xtmp[tl])
-        f2tmp,_,_ = f_avoid(xtmp)
+        if param.useBatch:
+            f2tmp,_,_ = f_avoid2(xtmp)
+        else:
+            f2tmp,_,_ = f_avoid(xtmp)
+            
         cost = np.linalg.norm(ftmp.flatten())**2 * param.Q_track + np.linalg.norm(f2tmp.flatten())**2 * param.Q_avoid + np.linalg.norm(utmp) * param.r
 
         if cost < cost0 or alpha < 1e-3:
@@ -218,8 +276,9 @@ for i in range( param.nbIter ):
         alpha /=2
     
     if np.linalg.norm(alpha * du) < 1e-2: # Early stop condition
-       break
-
+        break
+toc=time.time()
+print('Solving in {} seconds'.format(toc-tic))
 # Ploting
 # ===============================