code: def a2c_loss(state, action, reward, done, initial_core_state, model):
      # This time with gradients.
      (unused_actions, policy_logits, baseline), _ = model(
          state, done, initial_core_state, unroll=True
      )
  
      reward = reward[1:]
      done = done[1:]
      policy_logits = policy_logits[:-1]
      bootstrap_value = baseline[-1]
      baseline = baseline[:-1]
      action = action[:-1]
  
      T, B, *_ = reward.shape
  
      discount = ~done * DISCOUNT
      returns = torch.empty(T, B)
  
      acc = bootstrap_value.detach()
      for t in range(T - 1, -1, -1):
          acc = reward[t] + acc * discount[t]
          returns[t] = acc
  
      advantages = returns - baseline
  
      policy = F.softmax(policy_logits, dim=-1)
      log_policy = F.log_softmax(policy_logits, dim=-1)
  
      cross_entropy = torch.gather(
          log_policy, dim=-1, index=action[..., None]
      )
      cross_entropy.squeeze_(-1)
  
      # Alternative:
      # cross_entropy = F.nll_loss(
      #     F.log_softmax(torch.flatten(logits, 0, 1), dim=-1),
      #     target=torch.flatten(actions, 0, 1),
      #     reduction="none",
      # )
      # cross_entropy = cross_entropy.view_as(returns)
  
      pg_loss = -torch.mean(cross_entropy * advantages.detach())
      baseline_loss = 0.5 * torch.mean(advantages**2)
      entropy_loss = torch.mean(policy * log_policy)
  
      return pg_loss, baseline_loss, entropy_loss
  
file_path: moolib/examples/a2c.py
imports: import csv
  import os
  import time
  import gym
  import torch
  from torch import nn
  from torch.nn import functional
  import moolib
label_desc: Use image node_id 3 for calling a global function with example usage: a2c_loss(state, action, reward, done, initial_core_state, model) and returns: pg_loss,  baseline_loss,  entropy_loss