78f46822aec6cea5e216e1f7404d351835c54346,yarll/agents/sac.py,SAC,train,#SAC#,92

Before Change

        next_value_batch = self.target_value(state1_batch)
        softq_targets = reward_batch + (1 - terminal1_batch) * \
            self.config["gamma"] * tf.reshape(next_value_batch, [-1])
        softq_targets = tf.cast(tf.reshape(softq_targets, [self.config["batch_size"], 1]), tf.float32)

        actions, action_logprob = self.actor_network(state0_batch)
        new_softq = self.softq_network(state0_batch, actions)

        with tf.GradientTape() as softq_tape:
            softq = self.softq_network(state0_batch, action_batch)
            softq_loss = tf.reduce_mean(tf.square(softq - softq_targets))
        value_target = tf.stop_gradient(new_softq - action_logprob)
        with tf.GradientTape() as value_tape:
            values = self.value_network(state0_batch)
            value_loss = tf.reduce_mean(tf.square(values - value_target))
        advantage = tf.stop_gradient(action_logprob - new_softq + values)
        with tf.GradientTape() as actor_tape:
            _, action_logprob = self.actor_network(state0_batch)
            actor_loss = tf.reduce_mean(action_logprob * advantage)

        actor_gradients = actor_tape.gradient(actor_loss, self.actor_network.trainable_weights)
        softq_gradients = softq_tape.gradient(softq_loss, self.softq_network.trainable_weights)
        value_gradients = value_tape.gradient(value_loss, self.value_network.trainable_weights)

        self.actor_optimizer.apply_gradients(zip(actor_gradients, self.actor_network.trainable_weights))

After Change

        next_value_batch = self.target_value(state1_batch)
        softq_targets = reward_batch + (1 - terminal1_batch) * \
            self.config["gamma"] * tf.reshape(next_value_batch, [-1])
        softq_targets = tf.reshape(softq_targets, [self.config["batch_size"], 1])

        with tf.GradientTape() as softq_tape:
            softq = self.softq_network(state0_batch, action_batch)
            softq_loss = tf.reduce_mean(tf.square(softq - softq_targets))
        with tf.GradientTape() as value_tape:
            values = self.value_network(state0_batch)
        with tf.GradientTape() as actor_tape:
            actions, action_logprob = self.actor_network(state0_batch)
            new_softq = self.softq_network(state0_batch, actions)
            advantage = tf.stop_gradient(action_logprob - new_softq + values)
            actor_loss = tf.reduce_mean(action_logprob * advantage)
        value_target = tf.stop_gradient(new_softq - action_logprob)
        with value_tape:
            value_loss = tf.reduce_mean(tf.square(values - value_target))

        actor_gradients = actor_tape.gradient(actor_loss, self.actor_network.trainable_weights)
        softq_gradients = softq_tape.gradient(softq_loss, self.softq_network.trainable_weights)
        value_gradients = value_tape.gradient(value_loss, self.value_network.trainable_weights)

        self.actor_optimizer.apply_gradients(zip(actor_gradients, self.actor_network.trainable_weights))

Instances