[TensorFlow2] Critic Loss Calculation for actor_critic

[srihari-humbarwadi]

If I understand correctly, the code in tensorflow2/actor_critic.py implements the One-step Actor-Critic (episodic) algorithm given on page 332 of RLbook2020 by Sutton/barto (picture given below).

Here we can see that the critic parameters w are updated only using the gradient of the value function for the current state S
which is represented as grad(V(S, w)) in the pseudocode shown above. The update skips the gradient of the value function for the next state S'. This can again be seen in the pseudocode above, there is no grad(V(S', w)) present in the update rule for critic parameters w.

In the code given below, including state_value_, _ = self.actor_critic(state_) (L43) inside the GradientTape would result in grad(V(S', w)) appearing in the update for w, which contradicts the pseudocode shown above.

Youtube-Code-Repository/ReinforcementLearning/PolicyGradient/actor_critic/tensorflow2/actor_critic.py

Lines 40 to 45 in 1ef7605

	reward = tf.convert_to_tensor(reward, dtype=tf.float32) # not fed to NN
	with tf.GradientTape(persistent=True) as tape:
	state_value, probs = self.actor_critic(state)
	state_value_, _ = self.actor_critic(state_)
	state_value = tf.squeeze(state_value)
	state_value_ = tf.squeeze(state_value_)

Please let me know if there are some gaps in my understanding!