# A2C in TensorFlow 2 using model with two heads

I am implementing some of the basic reinforcement learning algorithms but ran into a problem with an online (one-step TD) A2C implementation where my reward seems to decrease over time instead of increasing (light blue in below picture).

I already implemented REINFORCE (offline, dark blue) and a Q-Actor-Critic (online with one-step TD, red) algorithm that both seem to work fine (no baselines, though).

The mean_episode_reward is calculated by running the policy 100 times in a simple, static 5x5 grid world with two terminal states that reward +1 or -1 and a discount of 0.99

I used a single-layer feed-forward neural network for REINFORCE and Q-AC for the policy and added a second head for A2C, now:

The training cycle is quite simple, where I get predictions for each action and the state-value from the model (for a state s), pick an action, collect a reward and get new action- and state-value-predictions. These are then used to calculate the advantage (one-step TD), i.e. $$advantage = reward + discount * v(s') - v(s)$$. This is then used to calculate the loss for the actor and critic which is then added (actor_loss + 0.5 * critic_loss) and fed to the gradient tape:

done = False

# initial
prediction, v_s = self.model(self.input_fn(), training=True)

while not done:
action = self.chooseAction(prediction) # [0, 1, 2, 3] for up, right, ...
reward, done = self.env.step(self.actions[int(action)])

prediction, v_s_prime = self.model(self.input_fn(), training=True)
advantage = reward + self.discount * v_s_prime - v_s