title: How to supervise latent variables in the variational Autoencoder (VAE)
body: <p>I'm trying to use the <strong>Variational Autoencoder</strong> for the <strong>lateral control</strong> of the vehicle. (End-to-End driving)<br/>
  The whole project is based on the paper of Alexander Amini:<br/>
  <em>"Variational Autoencoder for End-to-End Control of Autonomous Driving with Novelty Detection and Training De-biasing"</em> <br/>
  <a href="http://people.csail.mit.edu/rosman/papers/iros-2018-variational.pdf" rel="nofollow noreferrer">http://people.csail.mit.edu/rosman/papers/iros-2018-variational.pdf</a></p>
  
  <p>The idea is that I'm sending one semantic segmented image (from the center of the car) to the input, compress it to the certain vector of latent values (25x1), where one of the latent variables supposed to represent the value of curvature of the road. This value I'm using for the control of the car, the rest latent variables I need to be able to detect Novelty (images, where the NN was trained unsufficiently) and to unbias some new data sets (the VAE can be further trained). As part of this issue I'm only concerned about the first value (curvature of the road).</p>
  
  <p><strong>The problem is, that the VAE outputs wrong values (for the lateral control) after encoding part (before decoding)</strong></p>
  
  <p>My training set consists of images with corresponding labels (value of the curvature of the road)</p>
  
  <p>To train my VAE I'm using the loss function which consist of three parts:</p>
  
  <ol>
  <li>Supervising of the latent variable, this is L2 Norm between the latent variable (ŷ = z[:,0], means I'm taking the first latent variable from the latent vector) and the label.<br/>  <em>L_part_1(y, y_hat) = (y - y_hat)^2</em></li>
  <li>Reconstruction part, which is L1 Norm between original image and it's reconstructed version. <br/> <em>L_part_2(x,x_hat) = |x - x_hat|</em></li>
  <li>KL Divergence between the latent variables and a unit Gaussian, providing regularization for the latent space. <br/> 
  <em>L_part_3(z_mean, z_std_dev) = -0.5 (1-z_mean^2 - z_std_dev^2 + log(z_std_dev + 1e-10)^2)</em></li>
  </ol>
  
  <p>The total loss function has a following form: <br/> 
  <em>L = 0.033 x L_part_1 + 0.1 x L_part_2 + 0.001 x L_part_3</em> <br/> 
  Whose coefficients regularizing the importance of every loss part.<br/> </p>
  
  <p><em>Those loss functions deviate a bit from the trivial form of VAE loss functions. I took those from that paper, that mentioned above. I also tried the trivial one, bit I'll tell about it later.</em><br/> </p>
  
  <p>Here you can see the graphics of the loss function: <a href="https://ibb.co/x8GgVxq" rel="nofollow noreferrer">https://ibb.co/x8GgVxq</a> <br/> 
  And the result of the reconstructed images: <a href="https://ibb.co/D1p88PG" rel="nofollow noreferrer">https://ibb.co/D1p88PG</a></p>
  
  <p>As you can see, the first part of the loss function (Third on the image) looks weird. And when I use this trained VAE in inference mode, I'm getting nonsense results.</p>
  
  <p>I also tried different version of the loss function (trivial), where the first part (supervising one latent value) remains, the second part (reconstructed) has following form: <br/> 
  L_part_2(x,x_hat) = x * log(1e-10 + x_hat) + (1 - x) * log(1e-10 + 1 - x_hat)<br/> 
  and the third part (KL Divergence) looks like:<br/> 
  L_part_3(z_mean, z_std_dev) = -0.5 * (1 + z_std_dev - (z_mean)^2 - exp(z_std_dev))</p>
  
  <p>For that kind of configuration of the loss function I got following results: <br/> 
  [loss functions] <a href="https://ibb.co/VYV1F7F" rel="nofollow noreferrer">https://ibb.co/VYV1F7F</a> <br/> 
  [reconstructed images] <a href="https://ibb.co/ZTrpfxv" rel="nofollow noreferrer">https://ibb.co/ZTrpfxv</a> <br/> </p>
  
  <h2>Here is the key parts of my code with the first configuration of the loss function</h2>
  
  <pre><code>class VAE(tf.keras.Model):
      def __init__(self, latent_dim, counter =0):
          super(VAE, self).__init__()
          self.counter = counter
          self.latent_dim = latent_dim
          self.inference_net = tf.keras.Sequential([
              tf.keras.layers.InputLayer(input_shape=(66, 200, 3)),
              tf.keras.layers.Conv2D(
                  filters=24, kernel_size=5, strides=(2, 2), activation='relu'),
              tf.keras.layers.Conv2D(
                  filters=36, kernel_size=5, strides=(2, 2), activation='relu'),
              tf.keras.layers.Conv2D(
                  filters=48, kernel_size=5, strides=(2, 2), activation='relu'),
              tf.keras.layers.Conv2D(
                  filters=64, kernel_size=3, strides=(1, 1), activation='relu'),
              tf.keras.layers.Conv2D(
                  filters=64, kernel_size=3, strides=(1, 1), activation='relu'),
              tf.keras.layers.Flatten(),
              tf.keras.layers.Dense(units=1000, activation=tf.nn.relu),
              tf.keras.layers.Dense(units=100),
              tf.keras.layers.Dense(units=50)])
  
          self.generative_net = tf.keras.Sequential([
              tf.keras.layers.InputLayer(input_shape=(None,latent_dim)),
              tf.keras.layers.Dense(units=50),
              tf.keras.layers.Dense(units=100),
              tf.keras.layers.Dense(units=1000, activation=tf.nn.relu),
              tf.keras.layers.Dense(units=1152, activation=tf.nn.relu),
              tf.keras.layers.Reshape(target_shape=(1, 18, 64)),
              tf.keras.layers.Conv2DTranspose(
                  filters=64,kernel_size=3,strides=(1, 1),padding="VALID",activation='relu'),
              tf.keras.layers.Conv2DTranspose(
                  filters=48, kernel_size=3, strides=(1, 1), padding="VALID",activation='relu'),
              tf.keras.layers.Conv2DTranspose(
                  filters=36, kernel_size=5, strides=(2, 2), padding="VALID",activation='relu'),
              tf.keras.layers.Conv2DTranspose(
                  filters=24, kernel_size=5, strides=(2, 2), padding="VALID",activation='relu'),
              tf.keras.layers.Conv2DTranspose(
                  filters=3, kernel_size=5, strides=(2, 2), padding="VALID",activation='relu')
                                      ])
      def encode(self, x):
  
          mean, logvar = tf.split(self.inference_net(x), num_or_size_splits=2, axis=1)
          eps = tf.random.normal(tf.shape(mean))
          z = eps * tf.exp(logvar * .5) + mean
          return mean, logvar, z
  
  
      def decode(self, z, apply_sigmoid=False):
          logits = self.generative_net(z)
          if apply_sigmoid:
              probs = tf.sigmoid(logits)
              return probs
  
          return logits
  
      def compute_clothoid(self, image):
          test_im = np.reshape(image, (-1, 200, 66, 3))
          _, _, z = self.encode(test_im)
          clothoid = z[0][0]
          return clothoid
  
  optimizer = tf.keras.optimizers.Adam(1e-4)
  
  
  def compute_loss(model, x, labels):
      # print(labels)
      z_mean, z_std, z = model.encode(x)
      x_logit = model.decode(z)
      z_curv = z[:,0]
      x_sum = x * 255.
      x_logit_sum = x_logit * 255.
  
      with file_writer.as_default():
          tf.summary.image('input_image', x_sum, step=model.counter)
          tf.summary.image('output_image', x_logit_sum, step=model.counter)
  
      x_logit = tf.image.resize(x_logit, [66, 200])
      x = tf.reshape(
          x, [-1, image_dim[0] * image_dim[1] * image_dim[2]]
                      )
      x_logit = tf.reshape(
          x_logit, [-1, image_dim[0] * image_dim[1] * image_dim[2]]
                          )
  
      sup_lat_loss = tf.math.square(z_curv - labels)
  
      #reconstruction_loss = x * tf.math.log(1e-10 + x_logit) + (1 - x) * tf.math.log(1e-10 + 1 - x_logit)
      reconstruction_loss = tf.abs(x - x_logit)
      reconstruction_loss = tf.reduce_sum(reconstruction_loss, 1)
  
      # kl_div_loss =  1 - tf.square(z_mean) - tf.square(z_std)  + tf.math.log(tf.square(z_std + 1e-10))
      kl_div_loss =  1 + z_std - tf.square(z_mean) - tf.exp(z_std)
      kl_div_loss = -0.5 * tf.reduce_sum(kl_div_loss, 1)
  
      loss = tf.reduce_mean(reconstruction_loss + kl_div_loss)
  
      with file_writer.as_default():
          sup_lat_loss_show = c1 * tf.reduce_mean(sup_lat_loss)
          reconstruction_loss_show = c2 * tf.reduce_mean(reconstruction_loss)
          kl_div_loss_show = c3 * tf.reduce_mean(kl_div_loss)
  
          with tf.name_scope('loss_function'):
              tf.summary.scalar('sup_lat_loss', sup_lat_loss_show, step=model.counter)
              tf.summary.scalar('reconstruction_loss', reconstruction_loss_show, step=model.counter)
              tf.summary.scalar('kl_div_loss', kl_div_loss_show, step=model.counter)
  
          return loss
  
  </code></pre>
  
  <h2>Possible reasons for the issue</h2>
  
  <p>Maybe the way I assign <code>z_curv</code> is the reason, or the way I coded the loss function.</p>