text: from collections import OrderedDict
  
  import numpy as np
  import torch
  import torch.optim as optim
  from torch import nn as nn
  
  import rlkit.torch.pytorch_util as ptu
  from rlkit.core.eval_util import create_stats_ordered_dict
  from rlkit.torch.torch_rl_algorithm import TorchTrainer
  
  
  class R2GGNNTrainer(TorchTrainer):
      """
      Recursive Reasoning Graph; shared GNN
      """
      def __init__(
              self,
              env,
              context_graph,
              qf1,
              target_qf1,
              policy_n,
              cactor,
              qf2,
              target_qf2,
              deterministic_cactor_in_graph=True,
              deterministic_next_action=False,
              use_entropy_loss=True,
              use_entropy_reward=True,
              sum_n_loss=False, # use sum instead of mean for n agent losses
              use_cactor_entropy_loss=True,
              use_automatic_entropy_tuning=True,
              state_dependent_alpha=False,
              target_entropy=None,
              negative_sampling=False,
  
              discount=0.99,
              reward_scale=1.0,
  
              policy_learning_rate=1e-4,
              context_graph_learning_rate=1e-3,
              qf_learning_rate=1e-3, # not used
              qf_weight_decay=0.,
              init_alpha=1.,
              cactor_learning_rate=1e-4,
              target_hard_update_period=1000,
              tau=1e-2,
              use_soft_update=False,
              qf_criterion=None,
              pre_activation_weight=0.,
              optimizer_class=optim.Adam,
  
              min_q_value=-np.inf,
              max_q_value=np.inf,
  
              context_graph_optimizer=None,
              cactor_optimizer=None,
              policy_optimizer_n=None,
              alpha_optimizer_n=None,
              calpha_optimizer_n=None,
              log_alpha_n = None,
              log_calpha_n = None,
      ):
          super().__init__()
          self.env = env
          if qf_criterion is None:
              qf_criterion = nn.MSELoss()
          self.context_graph = context_graph
          self.qf1 = qf1
          self.target_qf1 = target_qf1
          self.qf2 = qf2
          self.target_qf2 = target_qf2
          self.policy_n = policy_n
          self.cactor = cactor
  
          self.deterministic_cactor_in_graph = deterministic_cactor_in_graph
          self.deterministic_next_action = deterministic_next_action
          self.sum_n_loss = sum_n_loss
          self.negative_sampling = negative_sampling
  
          self.discount = discount
          self.reward_scale = reward_scale
  
          self.policy_learning_rate = policy_learning_rate
          self.context_graph_learning_rate = context_graph_learning_rate
          self.qf_learning_rate = qf_learning_rate
          self.qf_weight_decay = qf_weight_decay
          self.cactor_learning_rate = cactor_learning_rate
          self.target_hard_update_period = target_hard_update_period
          self.tau = tau
          self.use_soft_update = use_soft_update
          self.qf_criterion = qf_criterion
          self.pre_activation_weight = pre_activation_weight
          self.min_q_value = min_q_value
          self.max_q_value = max_q_value
  
          if context_graph_optimizer:
              self.context_graph_optimizer = context_graph_optimizer
          else:
              self.context_graph_optimizer = optimizer_class(
                                      list(self.context_graph.parameters())\
                                      +list(self.qf1.parameters())\
                                      +list(self.qf2.parameters()),
                                      lr=self.context_graph_learning_rate,
                                   )
  
          if policy_optimizer_n:
              self.policy_optimizer_n = policy_optimizer_n
          else:
              self.policy_optimizer_n = [
                  optimizer_class(
                      self.policy_n[i].parameters(),
                      lr=self.policy_learning_rate,
                  ) for i in range(len(self.policy_n))]
          if cactor_optimizer:
              self.cactor_optimizer = cactor_optimizer
          else:
              self.cactor_optimizer = optimizer_class(
                                          self.cactor.parameters(),
                                          lr=self.cactor_learning_rate,
                                      )
  
          self.init_alpha = init_alpha
          self.use_entropy_loss = use_entropy_loss
          self.use_entropy_reward = use_entropy_reward
          self.use_cactor_entropy_loss = use_cactor_entropy_loss
          self.use_automatic_entropy_tuning = use_automatic_entropy_tuning
          self.state_dependent_alpha = state_dependent_alpha
          if self.use_automatic_entropy_tuning:
              if target_entropy:
                  self.target_entropy = target_entropy
              else:
                  self.target_entropy = -np.prod(self.env.action_space.shape).item()  # heuristic value from Tuomas
              if self.use_entropy_loss:
                  if log_alpha_n:
                      self.log_alpha_n = log_alpha_n
                  else:
                      self.log_alpha_n = [ptu.tensor([np.log(self.init_alpha)], requires_grad=True, dtype=torch.float32) for i in range(len(self.policy_n))]
                  if alpha_optimizer_n:
                      self.alpha_optimizer_n = alpha_optimizer_n
                  else:
                      if self.state_dependent_alpha:
                          self.alpha_optimizer_n = [
                                  optimizer_class(
                                      self.log_alpha_n[i].parameters(),
                                      lr=self.policy_learning_rate,
                                  ) for i in range(len(self.log_alpha_n))]
                      else:
                          self.alpha_optimizer_n = [
                              optimizer_class(
                                  [self.log_alpha_n[i]],
                                  lr=self.policy_learning_rate,
                              ) for i in range(len(self.log_alpha_n))]
  
              if self.use_cactor_entropy_loss:
                  if log_calpha_n:
                      self.log_calpha_n = log_calpha_n
                  else:
                      self.log_calpha_n = [ptu.tensor([np.log(self.init_alpha)], requires_grad=True, dtype=torch.float32) for i in range(len(self.policy_n))]
                  if calpha_optimizer_n:
                      self.calpha_optimizer_n = calpha_optimizer_n
                  else:
                      if self.state_dependent_alpha:
                          self.calpha_optimizer_n = [
                                  optimizer_class(
                                      self.log_calpha_n[i].parameters(),
                                      lr=self.policy_learning_rate,
                                  ) for i in range(len(self.log_calpha_n))]
                      else:
                          self.calpha_optimizer_n = [
                              optimizer_class(
                                  [self.log_calpha_n[i]],
                                  lr=self.policy_learning_rate,
                              ) for i in range(len(self.log_calpha_n))]
  
          self.eval_statistics = OrderedDict()
          self._n_train_steps_total = 0
          self._need_to_update_eval_statistics = True
  
      def train_from_torch(self, batch):
          rewards_n = batch['rewards'].detach()
          terminals_n = batch['terminals'].detach()
          obs_n = batch['observations'].detach()
          actions_n = batch['actions'].detach()
          next_obs_n = batch['next_observations'].detach()
  
          batch_size = rewards_n.shape[0]
          num_agent = rewards_n.shape[1]
          whole_obs = obs_n.view(batch_size, -1)
          whole_actions = actions_n.view(batch_size, -1)
          whole_next_obs = next_obs_n.view(batch_size, -1) 
  
          """
          Policy operations.
          """
          online_actions_n, online_pre_values_n, online_log_pis_n = [], [], []
          for agent in range(num_agent):
              policy_actions, info = self.policy_n[agent](
                  obs_n[:,agent,:], return_info=True,
              )
              online_actions_n.append(policy_actions)
              online_pre_values_n.append(info['preactivation'])
              online_log_pis_n.append(info['log_prob'])
          k0_actions = torch.stack(online_actions_n) # num_agent x batch x a_dim
          k0_actions = k0_actions.transpose(0,1).contiguous() # batch x num_agent x a_dim
  
          k0_inputs = torch.cat([obs_n, k0_actions],dim=-1)
          k0_contexts = self.context_graph(k0_inputs)  # batch x num_agent x c_dim
  
          k1_actions = self.cactor(k0_contexts, deterministic=self.deterministic_cactor_in_graph)
          k1_inputs = torch.cat([obs_n, k1_actions],dim=-1)
          k1_contexts = self.context_graph(k1_inputs)
  
          policy_gradients_n = []
          alpha_n = []
          for agent in range(num_agent):
              policy_actions = online_actions_n[agent]
              pre_value = online_pre_values_n[agent]
              log_pi = online_log_pis_n[agent]
              if self.pre_activation_weight > 0.:
                  pre_activation_policy_loss = (
                      (pre_value**2).sum(dim=1).mean()
                  )
              else:
                  pre_activation_policy_loss = torch.tensor(0.).to(ptu.device) 
              if self.use_entropy_loss:
                  if self.use_automatic_entropy_tuning:
                      if self.state_dependent_alpha:
                          alpha = self.log_alpha_n[agent](obs_n[:,agent,:]).exp()
                      else:
                          alpha = self.log_alpha_n[agent].exp()
                      alpha_loss = -(alpha * (log_pi + self.target_entropy).detach()).mean()
                      self.alpha_optimizer_n[agent].zero_grad()
                      alpha_loss.backward()
                      self.alpha_optimizer_n[agent].step()
                      if self.state_dependent_alpha:
                          alpha = self.log_alpha_n[agent](obs_n[:,agent,:]).exp().detach()
                      else:
                          alpha = self.log_alpha_n[agent].exp().detach()
                          alpha_n.append(alpha)
                  else:
                      alpha_loss = torch.tensor(0.).to(ptu.device)
                      alpha = torch.tensor(self.init_alpha).to(ptu.device)
                      alpha_n.append(alpha)
                  entropy_loss = (alpha*log_pi).mean()
              else:
                  entropy_loss = torch.tensor(0.).to(ptu.device)
  
              q_input = torch.cat([policy_actions,k1_contexts[:,agent,:]],dim=-1)
              q1_output = self.qf1(q_input)
              q2_output = self.qf2(q_input)
              q_output = torch.min(q1_output,q2_output)
              raw_policy_loss = -q_output.mean()
              policy_loss = (
                      raw_policy_loss +
                      pre_activation_policy_loss * self.pre_activation_weight +
                      entropy_loss
              )
  
              policy_gradients_n.append(torch.autograd.grad(policy_loss, self.policy_n[agent].parameters(),retain_graph=True))
  
              if self._need_to_update_eval_statistics:
                  self.eval_statistics['Policy Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      policy_loss
                  ))
                  self.eval_statistics['Raw Policy Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      raw_policy_loss
                  ))
                  self.eval_statistics['Preactivation Policy Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      pre_activation_policy_loss
                  ))
                  self.eval_statistics['Entropy Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      entropy_loss
                  ))
                  if self.use_entropy_loss:
                      if self.state_dependent_alpha:
                          self.eval_statistics.update(create_stats_ordered_dict(
                              'Alpha {}'.format(agent),
                              ptu.get_numpy(alpha),
                          ))
                      else:
                          self.eval_statistics['Alpha {} Mean'.format(agent)] = np.mean(ptu.get_numpy(
                              alpha
                          ))
                  self.eval_statistics.update(create_stats_ordered_dict(
                      'Policy Action {}'.format(agent),
                      ptu.get_numpy(policy_actions),
                  ))
  
          for agent in range(num_agent):
              # self.policy_optimizer_n[agent].zero_grad()
              for pid,p in enumerate(self.policy_n[agent].parameters()):
                  p.grad = policy_gradients_n[agent][pid]
              self.policy_optimizer_n[agent].step()
  
          """
          Critic operations.
          """
          with torch.no_grad():
              next_actions_n, next_log_pis_n = [], []
              for agent in range(num_agent):
                  next_actions, next_info = self.policy_n[agent](
                      next_obs_n[:,agent,:], return_info=True,
                      deterministic=self.deterministic_next_action,
                  )
                  next_actions_n.append(next_actions)
                  next_log_pis_n.append(next_info['log_prob'])
              next_k0_actions = torch.stack(next_actions_n) # num_agent x batch x a_dim
              next_k0_actions = next_k0_actions.transpose(0,1).contiguous() # batch x num_agent x a_dim
  
              next_k0_inputs = torch.cat([next_obs_n, next_k0_actions],dim=-1)
              next_k0_contexts = self.context_graph(next_k0_inputs)  # batch x num_agent x c_dim
  
              next_k1_actions = self.cactor(next_k0_contexts, deterministic=self.deterministic_cactor_in_graph)
              next_k1_inputs = torch.cat([next_obs_n, next_k1_actions],dim=-1)
              next_k1_contexts = self.context_graph(next_k1_inputs)
  
          buffer_inputs = torch.cat([obs_n, actions_n],dim=-1)
          buffer_contexts = self.context_graph(buffer_inputs) # batch x num_agent x c_dim
          q_inputs = torch.cat([actions_n, buffer_contexts],dim=-1)
          q1_preds_n = self.qf1(q_inputs)
          q2_preds_n = self.qf2(q_inputs)
  
          raw_qf1_loss_n, raw_qf2_loss_n, q_target_n = [], [], []
          for agent in range(num_agent):
              with torch.no_grad():
                  next_policy_actions = next_actions_n[agent]
                  next_log_pi = next_log_pis_n[agent]
                  next_q_input = torch.cat([next_policy_actions,next_k1_contexts[:,agent,:]],dim=-1)
  
                  next_target_q1_values = self.target_qf1(next_q_input)
                  next_target_q2_values = self.target_qf2(next_q_input)
                  next_target_q_values = torch.min(next_target_q1_values, next_target_q2_values)
  
                  if self.use_entropy_reward:
                      if self.state_dependent_alpha:
                          next_alpha = self.log_alpha_n[agent](next_obs_n[:,agent,:]).exp()
                      else:
                          next_alpha = alpha_n[agent]
                      next_target_q_values =  next_target_q_values - next_alpha * next_log_pi
  
                  q_target = self.reward_scale*rewards_n[:,agent,:] + (1. - terminals_n[:,agent,:]) * self.discount * next_target_q_values
                  q_target = torch.clamp(q_target, self.min_q_value, self.max_q_value)
                  q_target_n.append(q_target)
  
              q1_pred = q1_preds_n[:,agent,:]
              raw_qf1_loss = self.qf_criterion(q1_pred, q_target)
              raw_qf1_loss_n.append(raw_qf1_loss)
  
              q2_pred = q2_preds_n[:,agent,:]
              raw_qf2_loss = self.qf_criterion(q2_pred, q_target)
              raw_qf2_loss_n.append(raw_qf2_loss)
  
              if self._need_to_update_eval_statistics:
                  self.eval_statistics['QF1 Loss {}'.format(agent)] = np.mean(ptu.get_numpy(raw_qf1_loss))
                  self.eval_statistics['QF2 Loss {}'.format(agent)] = np.mean(ptu.get_numpy(raw_qf2_loss))
                  self.eval_statistics.update(create_stats_ordered_dict(
                      'Q1 Predictions {}'.format(agent),
                      ptu.get_numpy(q1_pred),
                  ))
                  self.eval_statistics.update(create_stats_ordered_dict(
                      'Q2 Predictions {}'.format(agent),
                      ptu.get_numpy(q2_pred),
                  ))
                  self.eval_statistics.update(create_stats_ordered_dict(
                      'Q Targets {}'.format(agent),
                      ptu.get_numpy(q_target),
                  ))
  
          if self.sum_n_loss:
              raw_qf1_loss = torch.sum(torch.stack(raw_qf1_loss_n))
              raw_qf2_loss = torch.sum(torch.stack(raw_qf2_loss_n))
          else:
              raw_qf1_loss = torch.mean(torch.stack(raw_qf1_loss_n))
              raw_qf2_loss = torch.mean(torch.stack(raw_qf2_loss_n))
  
          if self.negative_sampling:
              perturb_actions = actions_n.clone() # batch x agent x |A|
              batch_size, num_agent, a_dim = perturb_actions.shape
              perturb_agents = torch.randint(low=0,high=num_agent,size=(batch_size,))
              neg_actions = torch.rand(batch_size,a_dim)*2.-1. # ranged in -1 to 1
              perturb_actions[torch.arange(batch_size),perturb_agents,:] = neg_actions
                  
              perturb_inputs = torch.cat([obs_n,perturb_actions],dim=-1)
              perturb_contexts = self.context_graph(perturb_inputs) # batch x num_agent x c_dim
              perturb_q_inputs = torch.cat([actions_n, perturb_contexts],dim=-1)
              perturb_q1_preds = self.qf1(perturb_q_inputs)[torch.arange(batch_size),perturb_agents,:]
              perturb_q2_preds = self.qf2(perturb_q_inputs)[torch.arange(batch_size),perturb_agents,:]
              perturb_q_targets = torch.stack(q_target_n).transpose(0,1).contiguous()[torch.arange(batch_size),perturb_agents,:]
  
              neg_loss1 = self.qf_criterion(perturb_q1_preds, perturb_q_targets)
              neg_loss2 = self.qf_criterion(perturb_q2_preds, perturb_q_targets)
          else:
              neg_loss1, neg_loss2 = torch.tensor(0.).to(ptu.device), torch.tensor(0.).to(ptu.device)
  
          if self.qf_weight_decay > 0:
              reg_loss1 = self.qf_weight_decay * sum(
                  torch.sum(param ** 2)
                  for param in self.qf1.regularizable_parameters()
              )
  
              reg_loss2 = self.qf_weight_decay * sum(
                  torch.sum(param ** 2)
                  for param in self.qf2.regularizable_parameters()
              )
          else:
              reg_loss1, reg_loss2 = torch.tensor(0.).to(ptu.device), torch.tensor(0.).to(ptu.device)
  
          qf1_loss = raw_qf1_loss + reg_loss1 + neg_loss1
          qf2_loss = raw_qf2_loss + reg_loss2 + neg_loss2
  
          if self._need_to_update_eval_statistics:
              self.eval_statistics['raw_qf1_loss'] = np.mean(ptu.get_numpy(raw_qf1_loss))
              self.eval_statistics['raw_qf2_loss'] = np.mean(ptu.get_numpy(raw_qf2_loss))
              self.eval_statistics['neg_qf1_loss'] = np.mean(ptu.get_numpy(neg_loss1))
              self.eval_statistics['neg_qf2_loss'] = np.mean(ptu.get_numpy(neg_loss2))
              self.eval_statistics['reg_qf2_loss'] = np.mean(ptu.get_numpy(reg_loss1))
              self.eval_statistics['reg_qf2_loss'] = np.mean(ptu.get_numpy(reg_loss2))
  
          self.context_graph_optimizer.zero_grad()
          cg_loss = qf1_loss+qf2_loss
          cg_loss.backward()
          self.context_graph_optimizer.step()
  
          """
          Central actor operations.
          """
          buffer_inputs = torch.cat([obs_n, actions_n],dim=-1)
          buffer_contexts = self.context_graph(buffer_inputs) # batch x num_agent x c_dim
          cactor_loss_n = []
          for agent in range(num_agent):
              cactor_actions, cactor_info = self.cactor(
                  buffer_contexts[:,agent,:], return_info=True,
              )
              cactor_pre_value = cactor_info['preactivation']
              if self.pre_activation_weight > 0:
                  pre_activation_cactor_loss = (
                      (cactor_pre_value**2).sum(dim=1).mean()
                  )
              else:
                  pre_activation_cactor_loss = torch.tensor(0.).to(ptu.device)
              if self.use_cactor_entropy_loss:
                  cactor_log_pi = cactor_info['log_prob']
                  if self.use_automatic_entropy_tuning:
                      if self.state_dependent_alpha:
                          calpha = self.log_calpha_n[agent](whole_obs).exp()
                      else:
                          calpha = self.log_calpha_n[agent].exp()
                      calpha_loss = -(calpha * (cactor_log_pi + self.target_entropy).detach()).mean()
                      self.calpha_optimizer_n[agent].zero_grad()
                      calpha_loss.backward()
                      self.calpha_optimizer_n[agent].step()
                      if self.state_dependent_alpha:
                          calpha = self.log_calpha_n[agent](whole_obs).exp().detach()
                      else:
                          calpha = self.log_calpha_n[agent].exp().detach()
                  else:
                      calpha_loss = torch.tensor(0.).to(ptu.device)
                      calpha = torch.tensor(self.init_alpha).to(ptu.device)
                  cactor_entropy_loss = (calpha*cactor_log_pi).mean()
              else:
                  cactor_entropy_loss = torch.tensor(0.).to(ptu.device)
              
              q_input = torch.cat([cactor_actions,buffer_contexts[:,agent,:]],dim=-1)
              q1_output = self.qf1(q_input)
              q2_output = self.qf2(q_input)
              q_output = torch.min(q1_output,q2_output)
              raw_cactor_loss = -q_output.mean()
              cactor_loss = (
                      raw_cactor_loss +
                      pre_activation_cactor_loss * self.pre_activation_weight +
                      cactor_entropy_loss
              )
              cactor_loss_n.append(cactor_loss)
  
              if self._need_to_update_eval_statistics:
                  if self.use_cactor_entropy_loss:
                      if self.state_dependent_alpha:
                          self.eval_statistics.update(create_stats_ordered_dict(
                              'CAlpha {}'.format(agent),
                              ptu.get_numpy(calpha),
                          ))
                      else:
                          self.eval_statistics['CAlpha {} Mean'.format(agent)] = np.mean(ptu.get_numpy(
                              calpha
                          ))
                  self.eval_statistics['Cactor Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      cactor_loss
                  ))
                  self.eval_statistics['Raw Cactor Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      raw_cactor_loss
                  ))
                  self.eval_statistics['Preactivation Cactor Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      pre_activation_cactor_loss
                  ))
                  self.eval_statistics['Entropy Cactor Loss {}'.format(agent)] = np.mean(ptu.get_numpy(
                      cactor_entropy_loss
                  ))
          if self.sum_n_loss:
              cactor_loss = torch.sum(torch.stack(cactor_loss_n))
          else:
              cactor_loss = torch.mean(torch.stack(cactor_loss_n))
          self.cactor_optimizer.zero_grad()
          cactor_loss.backward()
          self.cactor_optimizer.step()
                  
          self._need_to_update_eval_statistics = False
          self._update_target_networks()
          self._n_train_steps_total += 1
  
      def _update_target_networks(self):
          if self.use_soft_update:
              ptu.soft_update_from_to(self.qf1, self.target_qf1, self.tau)
              ptu.soft_update_from_to(self.qf2, self.target_qf2, self.tau)
          else:
              if self._n_train_steps_total % self.target_hard_update_period == 0:
                  ptu.copy_model_params_from_to(self.qf1, self.target_qf1)
                  ptu.copy_model_params_from_to(self.qf2, self.target_qf2)
  
      def get_diagnostics(self):
          return self.eval_statistics
  
      def end_epoch(self, epoch):
          self._need_to_update_eval_statistics = True
  
      @property
      def networks(self):
          res = [
              *self.policy_n,
              self.context_graph,
              self.cactor,
              self.qf1,
              self.target_qf1,
              self.qf2,
              self.target_qf2,
          ]
          if self.state_dependent_alpha:
              res.extend(self.log_alpha_n)
              if self.use_cactor_entropy_loss:
                  res.extend(self.log_calpha_n)
          return res
  
      def get_snapshot(self):
          res = dict(
              context_graph=self.context_graph,
              qf1=self.qf1,
              target_qf1=self.target_qf1,
              qf2=self.qf2,
              target_qf2=self.target_qf2,
              cactor=self.cactor,
              policy_n=self.policy_n,
              # optimizers
              context_graph_optimizer=self.context_graph_optimizer,
              cactor_optimizer=self.cactor_optimizer,
              policy_optimizer_n=self.policy_optimizer_n,
          )
          if self.use_automatic_entropy_tuning:
              res['log_alpha_n'] = self.log_alpha_n
              res['alpha_optimizer_n'] = self.alpha_optimizer_n
              if self.use_cactor_entropy_loss:
                  res['log_calpha_n'] = self.log_calpha_n
                  res['calpha_optimizer_n'] = self.calpha_optimizer_n
          return res
  