How to set the target for the actor in A2C?
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.