Method _train_func

examples/reinforcement_learning/tutorial_C51.py:268–274

(self, b_o, b_index, b_m)

    @tf.function
    def _train_func(self, b_o, b_index, b_m):
        with tf.GradientTape() as tape:
            b_dist_a = tf.gather_nd(self.qnet(b_o), b_index)
            loss = tf.reduce_mean(tf.negative(tf.reduce_sum(b_dist_a * b_m, 1)))

        grad = tape.gradient(loss, self.qnet.trainable_weights)
        self.optimizer.apply_gradients(zip(grad, self.qnet.trainable_weights))


# #############################  Trainer  ###################################