How to set the target for the actor in A2C?

Question

I did a simple Actor-Critic implementation in Keras using 2 networks where the critic learns the Q-Values of every action, and the actor predicts probabilities for choosing each action. In training, the target probabilities for the actor was a one-hot vector with 1.0 in the maximum Q-Value prediction position and 0.0 in all the rest, and simply used fit method on the actor model with mean squared error loss function.

However, I'm not sure what to set as the target when switching to A2C. In all the guides I saw it's mentioned that the critic now learns one value per state, not one value per action in the action space.

This change makes it unclear on how to set the target vector for the actor. The guides/SE questions I went over did not explain this point and simply said that we can calculate the advantage value using the value function (here, here and here) for the current and next state, which is fine, except we can only do that for the specific action taken and not for every action in the action-space because we don't the value for every next state for every action.

In other words, we only know A(s,a) for our memorized a, and we know nothing about the advantage of other actions.

One of my guesses was that you still calculate the Q-Values, because after all, the value function is defined by the Q-Values. The value function is the sum over every action a of Q(s,a)*p(a). So does the critic need to learn the Q-Values and sum their multiplications with the probabilities generated by the policy network (actor), and calculate the advantages of every action?

It's even more confusing because in one of the guides they said that the critic actually learns the advantage values, and not the value function (like all the other guides said), which is strange because you need to use the critic to predict the value function of the state and the next state. Also, the advantage function is per-action and in the implementations I see the critic has one output neuron.

I think that what's being done in the examples I saw was to train the actor to fit a one-hot vector for the selected action (not the best action by the critic), but modify the loss-function value using the advantage value (possibly to influence the gradient). Is that the case?

Answer 1

In short, my last sentence was the correct answer. The "target" is a one-hot with the selected action, but there's a trick.

---

A2C Loss Function

A very crucial part of A2C implementation that I missed is the custom loss function that takes into account the advantage. The loss function multiplies the advantage with the negative log of the current probability to select the action that was selected.

The trick is that if the advantage is negative, the loss function will switch sign, so the gradients will be applied to the opposite direction.

In one dimension it's easier to understand. Let's say my target prediction is 1 and my actual prediction is 0.6. A simple loss would be defined as target - prediction, or in this case 0.4 and future predictions will be closer to one. If my prediction was 1.4, then the loss would be -0.4. A negative loss would mean predicting a lower result in the future, and a positive result would mean predicting a higher result in the future.

If the sign of the loss function is switched, the prediction will actually move away from 1.

The same thing happens when you multiply the advantage in the loss function. A negative advantage would mean that this action is worse than the value of the state so we need to avoid it, and a positive advantage means that the action is encouraged.

---

In Keras (Tensorflow 2.0):

Here's the custom loss function:

def custom_actor_loss(y_true, y_prediction, advantage):
    prediction = K.clip(y_prediction, 1e-8, 1 - 1e-8)
    log_probabilities = y_true * K.log(prediction)

    return K.sum(-log_probabilities*advantage)

The values are clipped because log of 0 is undefined.

And the rest of the network building:

input_layer = Input(shape=self._state_size, name='state_in')
advantage = Input(shape=[1], name='advantage')
target_prediction = Input(shape=self._actions_num, name='target')

inner_layer = Dense(units=layer_size, activation='relu')(input_layer)
actor_out = Dense(units=self._actions_num, activation='softmax', name='actor_out')(inner_layer)

self._actor = Model([input_layer, target_prediction, advantage], actor_out, name='actor')
self._actor.add_loss(custom_actor_loss(actor_out, target_prediction, advantage))
self._actor.compile(optimizer=Adam(learning_rate=actor_learning_rate))

And in the training loop (where future_rewards_prediction and critic_prediction are the outputs from the critic for the current and next state (except for the terminal state where the future_rewards_prediction is set to 0):

# Train actor
target_probabilities = np.zeros([1, self._actions_num])
target_probabilities[0][memory[step_idx].action] = 1.0
advantage = memory[step_idx].reward + future_rewards_prediction * self._future_discount - critic_prediction
self._actor.fit((memory[step_idx].state, target_probabilities, advantage), verbose=0)

*Notice how I don't really specify a y in my fit call. This is because of an issue I had when trying to implement a custom loss function in Keras which was solved by this answer.