Create Custom Loss Function to Minimize the Attitude Error - tensorflow

Question

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.
Why this loss function won't reduce the error correctly while mathematically it is correct?
How could I decrease the execution time of the loss function?
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

Answer 1

You should be able to vectorize (and speed up) the calculation using this approach. (I'm not sure I have all the signs right - don't understand why your lines 'x1 = -x1' are there. And I've temporarily remove the 'clip' part, you'll need to put that back if you want it)
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)))
Can't see what the error in the maths is, sorry