batch_loss_function.py

def transform_matrix_tensor(theta, d, a, alpha):
    # tensor version of transform matrix
    matrix = [[tf.cos(theta), tf.multiply(-tf.sin(theta), tf.cos(alpha)), tf.multiply(tf.sin(theta), tf.sin(alpha)), tf.multiply(a, tf.cos(theta))], 
              [tf.sin(theta), tf.multiply(tf.cos(theta), tf.cos(alpha)), tf.multiply(-tf.cos(theta), tf.sin(alpha)), tf.multiply(a, tf.sin(theta))], 
              [tf.zeros_like(theta), tf.sin(alpha), tf.cos(alpha), d], 
              [tf.zeros_like(theta), tf.zeros_like(theta), tf.zeros_like(theta), tf.ones_like(theta)]]
    return matrix
  
def batch_matmul(location_v, batch_theta_v):
    # perform matrix multiplication between the location vector and the transform matrix, 
    # independently for each example in the batch, but done in a parallel way
    zeros = tf.zeros_like(batch_theta_v)
    ones = tf.ones_like(batch_theta_v)
    m0 = transform_matrix_tensor(batch_theta_v, zeros, ones, zeros)
    m = tf.multiply(m0, location_v)
    m = tf.reduce_sum(m, axis=1)
    m = tf.transpose(m)
    return m
  
def forward_kinematics_loss_2(y_true, y_pred):
    # y_true is the xy position
    # y_pred is the 2-dimensional theta output
    theta1 = y_pred[:, 0]
    theta2 = y_pred[:, 1]
    zeros = tf.zeros_like(theta1)
    zeros = K.expand_dims(zeros, axis=1)
    
    location_v = K.concatenate([zeros, zeros, zeros, zeros+1], axis=1)
    location_v = K.expand_dims(location_v, axis=-1)
    location_v = K.concatenate([location_v]*4, axis=2)
    location_v = tf.transpose(location_v, perm=[2, 1, 0])
    
    end_tip_1st_segment = batch_matmul(location_v, theta1)
    
    location_v = K.expand_dims(end_tip_1st_segment, axis=-1)
    location_v = K.concatenate([location_v]*4, axis=2)
    location_v = tf.transpose(location_v, perm=[2, 1, 0])
    
    end_tip_2nd_segment = batch_matmul(location_v, theta2)
    
    xy = end_tip_2nd_segment[:, :2]
    loss1 = K.mean(K.square(xy - y_true))
    pi = tf.constant(math.pi)
    loss2 = tf.reduce_sum(tf.maximum(tf.abs(y_pred)-[[pi, 0.5 * pi]],0))
    loss = loss1 + loss2
    return loss