Create Custom Loss Function to Minimize the Attitude Error

I want to use IMU (accelerometer and gyroscope) readings to compute the attitude via Neural Network. The input will be input_shape = (time steps, 6) and the output is in the quaternion form output_shape = (time steps,4).

Based on mathematical calculations the error between reference quaternion and the predicted quaternion is

y_pred[i,]=w0,x0,y0,z0
y_true[i,]=w1,x1,y1,z1 
w = w0*w1 - x0*x1 - y0*y1 - z0*z1
x = w0*x1 + x0*w1 + y0*z1 - z0*y1
y = w0*y1 - x0*z1 + y0*w1 + z0*z1
z = w0*z1 + x0*y1 - y0*x1 + z0*w1
error_quaternion = [w, x, y, z]

To minimize the error, the scaler part of error quaternion (w) must be minimize. (please just ignore the math) So to reach the best predication the w must be minimized (w is the shortest angle between the predicted and reference attitude) -

Ref = {Markley, F. Landis, and John L. Crassidis. Fundamentals of spacecraft attitude determination and control. Vol. 1286. New York, NY, USA:: Springer New York, 2014.}

I write this loss function

def LossQuat2(y_true, y_pred):
        a, b = y_true.get_shape()
        error = []
        for i in range(a):
            w0,x0,y0,z0 = tf.unstack(y_pred[i,])
            w1,x1,y1,z1 = tf.unstack(y_true[i,])
            x1 = -x1
            y1 = -y1
            z1 = -z1
            w = w0*w1 - x0*x1 - y0*y1 - z0*z1
            error.append(2*tf.math.acos(K.clip(tf.math.sqrt(w*w), -1., 1.)))
        return tf.reduce_mean(error)

To validate it really calculate the error I try this code and the error calculated precisely

w0,x0,y0,z0 = y_pred[i,]
w1,x1,y1,z1 = y_true[i,]
x1 = -x1
y1 = -y1
z1 = -z1
w = w0*w1 - x0*x1 - y0*y1 - z0*z1
error = 2*math.acos(K.clip(np.sqrt(w*w), -1., 1.))

But after use this loss function to train the model, the output error is to much bigger than the MSE loss function. Also, it is too slow than MSE.

1. Why this loss function won't reduce the error correctly while mathematically it is correct?
2. How could I decrease the execution time of the loss function?
3. Is it true to use the for loop function? Is there any way to remove the for loop?

UPDATE:

Mathematics

Quaternion:

Quaternion is an attitude representation with 4 elements q=[w x y z]

w is the scalar part or real part

x y z are the vector part or imaginary part

Also, the quaternion could be written as:

q = [cos(theta/2) e*sin(theta/2)] , e is a unit vector (e=[i j k]

• I intend to estimate the quaternion by the neural network

Quaternion Inverse

quaternion inverse or quaternion conjugate can be calculated by:

quaternion = [w x y z]
inverse(quaternion) = [w -x -y -z]

Quaternion Multiplication

To find the difference between the estimated attitude and the true(reference) attitude, estimated attitude(NN output) must be multiply by the quaternion reference using quaternion multiplication.

Quaternion multiplication:

q_m = q1 * inverse(q2)

or

q_m = q2 * inverse(q1)

both of them is the same.

If

q1=w0,x0,y0,z0
q2=w1,x1,y1,z1 

Then q_m = [w x y z] and it could be calculated by:

w = w0*w1 - x0*x1 - y0*y1 - z0*z1
x = w0*x1 + x0*w1 + y0*z1 - z0*y1
y = w0*y1 - x0*z1 + y0*w1 + z0*z1
z = w0*z1 + x0*y1 - y0*x1 + z0*w1

The shortest angle between q1 and q2 is theta:

Theta = 2*acos(sqrt(w*w))

What I need is to write a loss function to minimize theta, if theta = 0, w will be equal to 1, so, the optimal q_m is:

q_m=[1 0 0 0]

Many thanks to david-harris @david-harris:

def loss(y_true, y_pred):
    z = y_true * y_pred * tf.constant([[1., -1., -1., -1.]])
    wtot = tf.reduce_sum(z, axis=1)
    return tf.reduce_mean(2*tf.math.acos(tf.math.sqrt(wtot*wtot)))

It is much faster but it seems that but it reduce all value of the quaternion, so it does not work correctly.

**

So sorry for lots of math.

**

UPDATE 2

Based on David's suggested code, I wrote this:

def loss(y_true, y_pred):
z = y_true * (y_pred * tf.constant([1., -1., -1., -1.000000000]))
wtot = tf.reduce_sum(z,1)
return tf.reduce_mean(2*tf.math.acos(K.clip(tf.math.sqrt(wtot*wtot), -1.,1.)))

This code reduced the loss but MSE grows exponentially. I understand that this code wont optimized for MSE directly, but due to mathematics the MSE also must decreases. After 10 epochs

loss: 0.0124 - mse: 227.4045 

One of outputs based on the custom loss

Orange = Reference

Blue = Estimated by NN

One of outputs based on the MSE loss function