Create transform matrix in tensor form for forward kinematics

transform_matrix_tensor.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 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]
    
    m0 = transform_matrix_tensor(theta1, 0.0, 1.0, 0.0)
    m1 = transform_matrix_tensor(theta2, 0.0, 1.0, 0.0)
    m2 = K.constant([0.0, 0.0, 0.0, 1.0], shape = (1, 4))
    
    m = tf.matmul(tf.matmul(m0, m1), m2)
    xy = m[:2]
    xy = tf.reshape(xy, (2,))
    loss1 = mean_squared_error(xy, y_true)
    pi = tf.constant(math.pi)
    loss2 = K.mean(tf.maximum(tf.abs(y_pred)-[[pi, 0.5 * pi]], 0))
    loss = loss1 + loss2
    return loss