title: Keras sequential model with altered loss does not learn
body: <p>(I am new to stackexchange, but I believe I have correctly classified this question. If there's something off about my question, please inform me.)</p>
  <p>I am trying to write a machine learning algorithm that learns to move an arm by contracting muscles. I have done my best to work out every possible bug I can think of but I have come to an impasse where every individual part of the program seems to run correctly, yet the algorithm does not learn. Fundamentally, all this model is doing is finding the inverse of a function by training a neural network to said function's inputs and outputs. The only thing that makes it even remotely nontrivial is that it uses an intermediary function when calculating the loss.</p>
  <p>Working in python with TensorFlow, we first define some constants and a function that converts deltoid and bicep muscle contractions to hand positions,</p>
  <pre><code>lH =1.0
  lU =1.0
  lCD=0.1
  lHD=0.1
  lHB=0.9
  lUB=0.1
  
  lD_max  = lCD+lHD
  lD_min  = abs(lCD-lHD)
  lD_diff = lD_max-lD_min
  lB_max  = lHB+lUB
  lB_min  = abs(lHB-lUB)
  lB_diff = lB_max-lB_min
  
  max_muscle_contraction = 0.9
  min_muscle_contraction = 0.1
  
  lD_min_eff = lD_min + min_muscle_contraction*lD_diff
  lD_max_eff = lD_min + max_muscle_contraction*lD_diff
  lB_min_eff = lB_min + min_muscle_contraction*lB_diff
  lB_max_eff = lB_min + max_muscle_contraction*lB_diff
  
  def contractionToPosition(c):
    # Takes a (n, m, 2) tensor of contraction pairs and returns a (n, m, 2) tensor of the resulting positions
    # Commonly takes (n, 2, 2) contraction tensors: a vector of initial and final vectors of deltoid-tricep pairs.
    cosD = (lCD**2+lHD**2 - tf.math.square(c[:,:,0]))/(2*lCD*lHD)
    cosD = tf.math.minimum(cosD, 2*max_muscle_contraction-1)
    cosD = tf.math.maximum(cosD, 2*min_muscle_contraction-1) # Equivalent to limiting the contraction
    sinD = tf.math.sqrt(1-tf.math.square(cosD))
    cosB = (lHB**2+lUB**2 - tf.math.square(c[:,:,1]))/(2*lHB*lUB)
    cosB = tf.math.minimum(cosB, 2*max_muscle_contraction-1)
    cosB = tf.math.maximum(cosB, 2*min_muscle_contraction-1) to limiting the contraction
    sinB = tf.math.sqrt(1-tf.math.square(cosB))
    px   = lH*cosD + lU*sinB*sinD - lU*cosB*cosD
    py   = -lH*sinD + lU*sinB*cosD + lU*cosB*sinD
    p    = tf.stack([px, py], axis=-1) # By px[i,j] being the [i,j]th px value that must be paired with the [i,j]th py value
    return p
  </code></pre>
  <p>Regardless of the above values and function's validity, the algorithm should still be able to learn from it because the data itself is synthetically generated with this same function. This function is also what the neural network is (approximately) trying to invert. Note that the neural network should take in the initial position and the planned final position, returning a change in the muscles contraction. Calculating the difference in the true final positions and the planned final positions will thus require that we also know the initial contraction. Toward this, we generate the synthetic data that we will later train the algorithm on,</p>
  <pre><code>def generateContraction(samples): # Returns a random vector of contraction lengths
    cD  = tf.zeros(samples)
    cD += tf.random.uniform(shape=cD.shape, minval=lD_min_eff, maxval=lD_max_eff)
    cB  = tf.zeros(samples)
    cB += tf.random.uniform(shape=cB.shape, minval=lB_min_eff, maxval=lB_max_eff)
    return tf.transpose(tf.stack([cD,cB]))
  
  def data(samples):
    ci = generateContraction(samples)
    cf = generateContraction(samples)
    c  = tf.stack([ci,cf], axis=1)
    p  = contractionToPosition(c)
    return p, c
  
  sample_size = 10000
  
  positions, contractions = data(sample_size)
  initial_contractions    = contractions[:,0]
  final_contractions      = contractions[:,1]
  
  features     = positions
  labels       = tf.subtract(final_contractions, initial_contractions)
  initial_data = initial_contractions
  </code></pre>
  <p>I have meticulously tested the entire process of this data's construction and every step has proven accurate. We then load this raw data into a dataset for the learning algorithm,</p>
  <pre><code>def load_array(data_arrays, batch_size, is_train=True):
    dataset = tf.data.Dataset.from_tensor_slices(data_arrays)
    if is_train:
      dataset = dataset.shuffle(buffer_size=1000)
    dataset = dataset.batch(batch_size)
    return dataset
  
  batch_size = 64
  data_iter = load_array((features, labels, initial_data), batch_size)
  </code></pre>
  <p>The network model doesn't need to be very complicated to know whether the learning algorithm works since there is no statistical error in the data. With this model, we are also intending it to act like the neural network found in the cerebellum of mammals. Specifically, this implies that for simple motions it is a shallow sequential neural network with ReLu activation. As such, we construct it fairly simply,</p>
  <pre><code>net = tf.keras.Sequential([
                             tf.keras.Input(shape=(2,2)),
                             tf.keras.layers.Flatten(),
                             tf.keras.layers.Dense(units=64,
                                                   activation='relu',
                                                   kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.1)),
                             tf.keras.layers.Dense(units=2,
                                                   activation='relu',
                                                   kernel_initializer=tf.keras.initializers.RandomNormal(stddev=0.1))
  ])
  </code></pre>
  <p>Finally, we write our learning algorithm based off the TensorFlow documentation, <a href="https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch" rel="nofollow noreferrer">https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch</a> . Note that we are optimizing the squared distance between the planned and actual final positions rather than optimizing the difference in the contractions. This is the only thing that makes this ever so slightly nontrivial.</p>
  <pre><code>loss = tf.keras.losses.MeanSquaredError()
  
  def train(net, train_iter, loss, epochs, lr):
    trainer = tf.keras.optimizers.Adam(learning_rate=lr)
    params = net.trainable_variables
    for epoch in range(epochs):
      epochError = 0
      for X, y, I in train_iter:
        with tf.GradientTape() as g:
          g.watch(params)
          P_hat = contractionToPosition(tf.reshape(net(X, training=True) + I, (-1,1,2)))
          P     = contractionToPosition(tf.reshape(                    y + I, (-1,1,2))) # We have to reshape because of our function contractionToPosition
          l = loss(P, P_hat)
        epochError += l
        error = l
        grads = g.gradient(l, params)
        trainer.apply_gradients(zip(grads,params))
      print(f'epoch {epoch + 1}, '
            f'loss: {epochError}')
  
  train(net, data_iter, loss, 5, 0.05)
  </code></pre>
  <p>The result of all this, though, is a complete lack of learning. Usually the epoch loss is about 109 (which is expected for no learning) with no significant change in said loss (usually fluctuates within +/-0.7 .) If anything is at fault, I would suspect this final code snippet to be, specifically the gradient tape. I have probed every aspect of the gradient tape, however, and everything seems to be functioning correctly. Overall, I cannot think of a part of my code I have not dissected at this point so I am at a total loss here.</p>
  <p>Any and all help is deeply appreciated!</p>