title: Implementing A3C on TensorFlow 2
body: <p>After finishing Coursera's Practical RL course on A3C, I'm trying to implement my own A3C agent using tensorflow 2. To start, I'm training it on the Cartpole environment but I can't get good results. For now, I've already launched several training with the following code, changing the entropy coefficient to see its impact (the results are shown below). Does it come from my implementation, or is it more a fine-tuning issue ?</p>
  <pre><code>class A3C:
    def __init__(self, state_dim, n_actions, optimizer=tf.keras.optimizers.Adam(1e-3)):
      self.state_input = Input(shape=state_dim)
      self.x = Dense(256, activation='relu')(self.state_input)
  
      self.head_v = Dense(1, activation='linear')(self.x)
      self.head_p = Dense(n_actions, activation='linear')(self.x)
  
      self.network = tf.keras.Model(inputs=[self.state_input], outputs=[self.head_v, self.head_p])
      self.optimizer = optimizer
  
    def forward(self, state):
      return self.network(state)
  
    def sample(self, logits):
      policy = np.exp(logits.numpy()) / np.sum(np.exp(logits.numpy()), axis=-1, keepdims=True)
      return np.array([np.random.choice(len(p), p=p) for p in policy])
  
  def evaluate(agent, env, n_games=1): &quot;&quot;&quot;Plays an a game from start till done, returns per-game rewards &quot;&quot;&quot;
    game_rewards = []
    for _ in range(n_games):
      state = env.reset()
  
      total_reward = 0
      while True:
        action = agent.sample(agent.forward(np.array([state]))[1])[0]
        state, reward, done, info = env.step(action)
        total_reward += reward
        if done: break
  
      game_rewards.append(total_reward)
    return game_rewards
  
  class EnvBatch:
    def __init__(self, n_envs = 10):
      self.envs = [gym.make(env_id) for _ in range(n_envs)]
      
    def reset(self):
      return np.array([env.reset() for env in self.envs])
  
    def step(self, actions):
      results = [env.step(a) for env, a in zip(self.envs, actions)]
      new_obs, rewards, done, infos = map(np.array, zip(*results))
      
      for i in range(len(self.envs)):
        if done[i]:
          new_obs[i] = self.envs[i].reset()
      
      return new_obs, rewards, done, infos
  
   
  
  env_id = &quot;CartPole-v0&quot;
  env = gym.make(env_id)
  state_dim = env.observation_space.shape
  n_actions = env.action_space.n
  agent = A3C(state_dim, n_actions)    
  env_batch = EnvBatch(10)
  batch_states = env_batch.reset()
  gamma=0.99
  
  rewards_history = []
  entropy_history = []
  
  for i in trange(200000):
    with tf.GradientTape() as t:
      batch_values, batch_logits = agent.forward(batch_states)
      batch_actions = agent.sample(batch_logits)
      batch_next_states, batch_rewards, batch_dones, _ = env_batch.step(batch_actions)
      batch_next_values, btach_next_logits = agent.forward(batch_next_states)
      batch_next_values *= (1 - batch_dones)
  
      probs = tf.nn.softmax(batch_logits)
      logprobs = tf.nn.log_softmax(batch_logits)
  
      logp_actions = tf.reduce_sum(logprobs * tf.one_hot(batch_actions, n_actions), axis=-1)
  
      advantage = batch_rewards + gamma*batch_next_values - batch_values
      entropy = -tf.reduce_sum(probs * logprobs, 1, name=&quot;entropy&quot;)
  
      actor_loss =  - tf.reduce_mean(logp_actions * tf.stop_gradient(advantage)) - 0.005 * tf.reduce_mean(entropy)
  
      target_state_values = batch_rewards + gamma*batch_next_values
      critic_loss = tf.reduce_mean((batch_values - tf.stop_gradient(target_state_values))**2 )
  
      loss = actor_loss + critic_loss
  
    var_list = agent.network.trainable_variables
    grads = t.gradient(loss,var_list)
    agent.optimizer.apply_gradients(zip(grads, var_list))
    batch_states = batch_next_states
  
    entropy_history.append(np.mean(entropy))
  
    if i % 500 == 0:
      if i % 2500 == 0:
        rewards_history.append(np.mean(evaluate(agent, env, n_games=3)))
  
      clear_output(True)
      plt.figure(figsize=[8, 4])
      plt.subplot(1, 2, 1)
      plt.plot(rewards_history, label='rewards')
      plt.title(&quot;Session rewards&quot;)
      plt.grid()
      plt.legend()
  
      plt.subplot(1, 2, 2)
      plt.plot(entropy_history, label='entropy')
      plt.title(&quot;Policy entropy&quot;)
      plt.grid()
      plt.legend()
      plt.show()
  </code></pre>
  <ol>
  <li><a href="https://i.stack.imgur.com/mIxQO.png" rel="nofollow noreferrer">Beta = 0.005 - Training 1</a></li>
  <li><a href="https://i.stack.imgur.com/yBxJ4.png" rel="nofollow noreferrer">Beta = 0.005 - Training 2</a></li>
  <li><a href="https://i.stack.imgur.com/aW657.png" rel="nofollow noreferrer">Beta = 0.005 - Training 3</a></li>
  <li><a href="https://i.stack.imgur.com/HZN8i.png" rel="nofollow noreferrer">Beta = 0.05 - Training 1</a></li>
  <li><a href="https://i.stack.imgur.com/GQ0p2.png" rel="nofollow noreferrer">Beta = 0.05 - Training 2</a></li>
  <li><a href="https://i.stack.imgur.com/vrd8U.png" rel="nofollow noreferrer">Beta = 0.05 - Training 3</a></li>
  </ol>