UnboundLocalError: 'mse_loss' is used before assignment when using `value: classical` in yaml file

[kennyfrc]

Hi There,

I'm doing some supervised learning with the use of engine scores as policy data. Because of that, I'm using value: classical in my yaml file.

Once I begin the training, I get this error in my tfprocess.py file:

UnboundLocalError: 'mse_loss' is used before assignment

This is triggered by the code below. I noticed that mse_loss is not defined if I choose a value that is not wdl.

    def process_inner_loop(self, x, y, z, q, m):
        with tf.GradientTape() as tape:
            outputs = self.model(x, training=True)
            policy = outputs[0]
            value = outputs[1]
            policy_loss = self.policy_loss_fn(y, policy)
            reg_term = sum(self.model.losses)
            if self.wdl:
                value_ce_loss = self.value_loss_fn(self.qMix(z, q), value)
                value_loss = value_ce_loss
            else:
                value_mse_loss = self.mse_loss_fn(self.qMix(z, q), value)
                value_loss = value_mse_loss
            if self.moves_left:
                moves_left = outputs[2]
                moves_left_loss = self.moves_left_loss_fn(m, moves_left)
            else:
                moves_left_loss = tf.constant(0.)

            total_loss = self.lossMix(policy_loss, value_loss,
                                      moves_left_loss) + reg_term
            if self.loss_scale != 1:
                total_loss = self.optimizer.get_scaled_loss(total_loss)
        if self.wdl:
            mse_loss = self.mse_loss_fn(self.qMix(z, q),  #value)
        else:
            value_loss = self.value_loss_fn(self.qMix(z, q), value)
        return policy_loss, value_loss, mse_loss, moves_left_loss, reg_term, tape.gradient(
            total_loss, self.model.trainable_weights)

To patch this, I simply assigned mse_loss = value_mse_loss like below:

        if self.wdl:
            mse_loss = self.mse_loss_fn(self.qMix(z, q), value)
        else:
            mse_loss = value_mse_loss
            value_loss = self.value_loss_fn(self.qMix(z, q), value)

After this, the code has worked. I'm not exactly very knowledgable about machine learning so let me know if the above makes sense. If it does, I'll submit a pull request.