text: from rl_games.common import a2c_common
  from rl_games.algos_torch import torch_ext
  
  from rl_games.algos_torch import central_value
  from rl_games.common import common_losses
  from rl_games.common import datasets
  from rl_games.algos_torch import ppg_aux
  from rl_games.common.ewma_model import EwmaModel
  
  from torch import optim
  import torch 
  from torch import nn
  import numpy as np
  import gym
  
  class A2CAgent(a2c_common.ContinuousA2CBase):
      def __init__(self, base_name, params):
          a2c_common.ContinuousA2CBase.__init__(self, base_name, params)
          obs_shape = self.obs_shape
          build_config = {
              'actions_num' : self.actions_num,
              'input_shape' : obs_shape,
              'num_seqs' : self.num_actors * self.num_agents,
              'value_size': self.env_info.get('value_size',1),
              'normalize_value' : self.normalize_value,
              'normalize_input': self.normalize_input,
          }
          
          self.model = self.network.build(build_config)
          self.model.to(self.ppo_device)
          self.states = None
          if self.ewma_ppo:
              self.ewma_model = EwmaModel(self.model, ewma_decay=0.889)
          self.init_rnn_from_model(self.model)
          self.last_lr = float(self.last_lr)
          self.bound_loss_type = self.config.get('bound_loss_type', 'bound') # 'regularisation' or 'bound'
          self.optimizer = optim.Adam(self.model.parameters(), float(self.last_lr), eps=1e-08, weight_decay=self.weight_decay)
  
          if self.has_central_value:
              cv_config = {
                  'state_shape' : self.state_shape, 
                  'value_size' : self.value_size,
                  'ppo_device' : self.ppo_device, 
                  'num_agents' : self.num_agents, 
                  'horizon_length' : self.horizon_length,
                  'num_actors' : self.num_actors, 
                  'num_actions' : self.actions_num, 
                  'seq_len' : self.seq_len,
                  'normalize_value' : self.normalize_value,
                  'network' : self.central_value_config['network'],
                  'config' : self.central_value_config, 
                  'writter' : self.writer,
                  'max_epochs' : self.max_epochs,
                  'multi_gpu' : self.multi_gpu
              }
              self.central_value_net = central_value.CentralValueTrain(**cv_config).to(self.ppo_device)
  
          self.use_experimental_cv = self.config.get('use_experimental_cv', True)
          self.dataset = datasets.PPODataset(self.batch_size, self.minibatch_size, self.is_discrete, self.is_rnn, self.ppo_device, self.seq_len)
          if self.normalize_value:
              self.value_mean_std = self.central_value_net.model.value_mean_std if self.has_central_value else self.model.value_mean_std
          if 'phasic_policy_gradients' in self.config:
              self.has_phasic_policy_gradients = True
              self.ppg_aux_loss = ppg_aux.PPGAux(self, self.config['phasic_policy_gradients'])
          self.has_value_loss = (self.has_central_value and self.use_experimental_cv) \
                              or (not self.has_phasic_policy_gradients and not self.has_central_value) 
          self.algo_observer.after_init(self)
      def update_epoch(self):
          self.epoch_num += 1
          return self.epoch_num
          
      def save(self, fn):
          state = self.get_full_state_weights()
          torch_ext.save_checkpoint(fn, state)
  
      def restore(self, fn):
          checkpoint = torch_ext.load_checkpoint(fn)
          self.set_full_state_weights(checkpoint)
  
      def get_masked_action_values(self, obs, action_masks):
          assert False
  
      def calc_gradients(self, input_dict):
          value_preds_batch = input_dict['old_values']
          old_action_log_probs_batch = input_dict['old_logp_actions']
          advantage = input_dict['advantages']
          old_mu_batch = input_dict['mu']
          old_sigma_batch = input_dict['sigma']
          return_batch = input_dict['returns']
          actions_batch = input_dict['actions']
  
          # if self.modality_dict is not None:
          #     obs_batch = {}
          #     for modality_name in self.modality_dict:
          #         obs_batch[modality_name] = input_dict[modality_name]
          #     obs_batch = self._preproc_obs(obs_batch)
          #
          # else:
          obs_batch = input_dict['obs']
          obs_batch = self._preproc_obs(obs_batch)
  
          lr_mul = 1.0
          curr_e_clip = lr_mul * self.e_clip
  
          batch_dict = {
              'is_train': True,
              'prev_actions': actions_batch, 
              'obs' : obs_batch,
          }
  
          rnn_masks = None
          if self.is_rnn:
              rnn_masks = input_dict['rnn_masks']
              batch_dict['rnn_states'] = input_dict['rnn_states']
              batch_dict['seq_length'] = self.seq_len
              
          with torch.cuda.amp.autocast(enabled=self.mixed_precision):
              res_dict = self.model(batch_dict)
              action_log_probs = res_dict['prev_neglogp']
              values = res_dict['values']
              entropy = res_dict['entropy']
              mu = res_dict['mus']
              sigma = res_dict['sigmas']
  
              if self.ewma_ppo:
                  ewma_dict = self.ewma_model(batch_dict)
                  proxy_neglogp = ewma_dict['prev_neglogp']
                  a_loss = common_losses.decoupled_actor_loss(old_action_log_probs_batch, action_log_probs, proxy_neglogp, advantage, curr_e_clip)
                  old_action_log_probs_batch = proxy_neglogp # to get right statistic later
                  old_mu_batch = ewma_dict['mus']
                  old_sigma_batch = ewma_dict['sigmas']
              else:
                  a_loss = common_losses.actor_loss(old_action_log_probs_batch, action_log_probs, advantage, self.ppo, curr_e_clip)
  
              if self.has_value_loss:
                  c_loss = common_losses.critic_loss(value_preds_batch, values, curr_e_clip, return_batch, self.clip_value)
              else:
                  c_loss = torch.zeros(1, device=self.ppo_device)
              if self.bound_loss_type == 'regularisation':
                  b_loss = self.reg_loss(mu)
              elif self.bound_loss_type == 'bound':
                  b_loss = self.bound_loss(mu)
              else:
                  b_loss = torch.zeros(1, device=self.ppo_device)
              losses, sum_mask = torch_ext.apply_masks([a_loss.unsqueeze(1), c_loss, entropy.unsqueeze(1), b_loss.unsqueeze(1)], rnn_masks)
              a_loss, c_loss, entropy, b_loss = losses[0], losses[1], losses[2], losses[3]
  
              loss = a_loss + 0.5 * c_loss * self.critic_coef - entropy * self.entropy_coef + b_loss * self.bounds_loss_coef
              
              if self.multi_gpu:
                  self.optimizer.zero_grad()
              else:
                  for param in self.model.parameters():
                      param.grad = None
  
          self.scaler.scale(loss).backward()
          #TODO: Refactor this ugliest code of they year
          if self.truncate_grads:
              if self.multi_gpu:
                  self.optimizer.synchronize()
                  self.scaler.unscale_(self.optimizer)
                  nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm)
                  with self.optimizer.skip_synchronize():
                      self.scaler.step(self.optimizer)
                      self.scaler.update()
              else:
                  self.scaler.unscale_(self.optimizer)
                  nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm)
                  self.scaler.step(self.optimizer)
                  self.scaler.update()    
          else:
              self.scaler.step(self.optimizer)
              self.scaler.update()
  
          with torch.no_grad():
              reduce_kl = rnn_masks is None
              kl_dist = torch_ext.policy_kl(mu.detach(), sigma.detach(), old_mu_batch, old_sigma_batch, reduce_kl)
              if rnn_masks is not None:
                  kl_dist = (kl_dist * rnn_masks).sum() / rnn_masks.numel()  #/ sum_mask
  
          if self.ewma_ppo:
              self.ewma_model.update()                    
  
          self.diagnostics.mini_batch(self,
          {
              'values' : value_preds_batch,
              'returns' : return_batch,
              'new_neglogp' : action_log_probs,
              'old_neglogp' : old_action_log_probs_batch,
              'masks' : rnn_masks
          }, curr_e_clip, 0)      
  
          self.train_result = (a_loss, c_loss, entropy, \
              kl_dist, self.last_lr, lr_mul, \
              mu.detach(), sigma.detach(), b_loss)
  
      def train_actor_critic(self, input_dict):
          self.calc_gradients(input_dict)
          return self.train_result
  
      def reg_loss(self, mu):
          if self.bounds_loss_coef is not None:
              reg_loss = (mu*mu).sum(axis=-1)
          else:
              reg_loss = 0
          return reg_loss
  
      def bound_loss(self, mu):
          if self.bounds_loss_coef is not None:
              soft_bound = 1.1
              mu_loss_high = torch.clamp_min(mu - soft_bound, 0.0)**2
              mu_loss_low = torch.clamp_max(mu + soft_bound, 0.0)**2
              b_loss = (mu_loss_low + mu_loss_high).sum(axis=-1)
          else:
              b_loss = 0
          return b_loss
  
  
  