text: import numpy as np
  import random
  
  import torch
  import torch.optim as optim
  import torch.nn.functional as F
  
  from nn import Actor, Critic
  from utils import OUNoise, ReplayBuffer
  
  class DDPG_Agent():
      """Deep Deterministic Policy Gradients (DDPG) Actor-Critic Agent."""
      
      def __init__(self, state_size, action_size, n_agents, random_seed, config):
          """Initialize Agent object.
          
          Params
          ======
              state_size (int): dimension of each state.
              action_size (int): dimension of each action.
              n_agents (int): number of agents (for multi-agent training).
              random_seed (int): random seed.
              config (obj): Config class object with model configuration.
          """
          self.config = config
          self.state_size = state_size
          self.action_size = action_size
          self.seed = random.seed(random_seed)
          self.n_agents = n_agents
  
          # Actor Network (w/ Target Network)
          self.actor_local = Actor(state_size, action_size, random_seed, 
                                   config.fc1_act, config.fc2_act, config.bn).to(config.DEVICE)
          self.actor_target = Actor(state_size, action_size, random_seed, 
                                    config.fc1_act, config.fc2_act, config.bn).to(config.DEVICE)
          self.actor_optimizer = optim.Adam(self.actor_local.parameters(), lr=config.lr_actor)
  
          # Critic Network (w/ Target Network)
          self.critic_local = Critic(state_size, action_size, n_agents, random_seed, 
                                     config.fc1_crit, config.fc2_crit, config.fc3_crit, config.bn).to(config.DEVICE)
          self.critic_target = Critic(state_size, action_size, n_agents, random_seed, 
                                      config.fc1_crit, config.fc2_crit, config.fc3_crit, config.bn).to(config.DEVICE)
          self.critic_optimizer = optim.Adam(self.critic_local.parameters(), lr=config.lr_critic, 
                                             weight_decay=config.l2_reg)
          
          # make sure local and target networks start with same weights
          self.hard_update(self.actor_local, self.actor_target)
          self.hard_update(self.critic_local, self.critic_target)
  
          # Noise process
          self.noise = OUNoise(action_size, random_seed)
          self.noise_decay = config.noise_decay
  
      def act(self, state, add_noise=True):
          """Returns actions for given state as per current policy."""
          state = torch.from_numpy(state).float().to(self.config.DEVICE)
          self.actor_local.eval()
          with torch.no_grad():
              action = self.actor_local(state).cpu().data.numpy()
          self.actor_local.train()
          if add_noise:
              self.noise_decay = np.maximum(self.noise_decay*self.config.noise_decay_factor, self.config.min_noise)
              action += self.noise_decay * self.noise.sample()
          return np.clip(action, -1, 1)
      
      def target_act(self, state):
          """Returns actions for given state as per target policy."""
          action = self.actor_target(state).cpu().data.numpy()
          return np.clip(action, -1, 1)
  
      def reset(self):
          self.noise.reset()
          
      def hard_update(self, local_model, target_model):
          """Hard update model parameters.
          θ_target = θ_local
          Params
          ======
              local_model: PyTorch model (weights will be copied from)
              target_model: PyTorch model (weights will be copied to)
          """
          for target_params, local_param in zip(target_model.parameters(),local_model.parameters()):
              target_params.data.copy_(local_param.data)   
              
  class MADDPG:
      """Multi-Agent Deep Deterministic Policy Gradients (MADDPG) Actor-Critic Agent."""
      def __init__(self, state_size, action_size, n_agents, config, random_seed=420):
          """Initialize Multi-Agent object.
          
          Params
          ======
              state_size (int): dimension of each state.
              action_size (int): dimension of each action.
              n_agents (int): number of agents (for multi-agent training).
              config (obj): Config class object with model configuration.
              random_seed (int): random seed.
          """
          self.config = config
          self.maddpg_agent = [DDPG_Agent(state_size, action_size, n_agents, random_seed, config) 
                              for _ in range(n_agents)]
          self.n_agents = n_agents
          
          # Replay memory
          self.memory = ReplayBuffer(config.buffer_size, config.batch_size, n_agents, random_seed)
          
          # Initialize time step (for updating every UPDATE_EVERY steps)
          self.t_step = 0
  
      def tensor(self, np_arr):
          """Numpy -> Pytorch tensor"""
          return torch.from_numpy(np_arr).float().to(self.config.DEVICE)
      
      def reset(self):
          """Reset noise of all agents."""
          for ddpg_agent in self.maddpg_agent:
              ddpg_agent.reset()
      
      def get_noise(self):
          """Get noise of all agents."""
          return [ddpg_agent.noise_decay for ddpg_agent in self.maddpg_agent]
      
      def act(self, obs_all_agents, add_noise=True):
          """Get actions from all agents in the MADDPG object."""
          actions = [agent.act(obs, add_noise) for agent, obs in zip(self.maddpg_agent, obs_all_agents)]
          return actions
  
      def target_act(self, obs_all_agents):
          """Get target network actions from all the agents in the MADDPG object."""
          target_actions = [self.tensor(ddpg_agent.target_act(obs)) for ddpg_agent, obs in 
                            zip(self.maddpg_agent, obs_all_agents)]
          return target_actions
      
      def step(self, state, action, reward, next_state, done):
          """Save experience in replay memory, and use random sample from buffer to learn."""
          # Save experience / reward
          self.memory.add(state, action, reward, next_state, done)
  
          # Learn, if enough samples are available in memory
          self.t_step = (self.t_step + 1) % self.config.update_every
  
          if len(self.memory) > self.config.batch_size and self.t_step == 0:
              for _ in range(self.config.update_times):
                  for a_i, _ in enumerate(self.maddpg_agent):
                      experiences = self.memory.sample()
                      self.learn(experiences, a_i)
                      
      def learn(self, experiences, agent_number):
          """Update policy and value parameters using given batch of experience tuples.
          Q_targets = r + γ * critic_target(next_state, actor_target(next_state))
          where:
              actor_target(state) -> action
              critic_target(state, action) -> Q-value
  
          Params
          ======
              experiences (Tuple[torch.Tensor]): tuple of (s, a, r, s', done) tuples.
              agent_number (int): agent number (index).
          """
          config = self.config
          
          agent = self.maddpg_agent[agent_number]
  
          states, actions, rewards, next_states, dones = experiences
          all_states = torch.cat(states, dim=1).to(self.config.DEVICE)
          all_next_states = torch.cat(next_states, dim=1).to(self.config.DEVICE)
  
          # ---------------------------- update critic ---------------------------- #
          # Get predicted next-state actions and Q values from target models
          all_next_action = self.target_act(next_states)
          all_next_action = torch.cat(all_next_action, dim=1).to(self.config.DEVICE)
          
          with torch.no_grad():
              Q_targets_next = agent.critic_target(all_next_states, all_next_action)
  
          # Compute Q targets for current states (y_i)
          Q_targets = \
              rewards[agent_number] + (config.gamma * Q_targets_next * (1 - dones[agent_number]))
  
          # Compute critic loss
          all_actions = torch.cat(actions, dim=1).to(self.config.DEVICE)
          Q_expected = agent.critic_local(all_states, all_actions)
          critic_loss = F.mse_loss(Q_expected, Q_targets.detach())
  
          # Minimize the loss
          agent.critic_optimizer.zero_grad()
          critic_loss.backward()
          if config.clip_grad_crit:
              torch.nn.utils.clip_grad_norm_(agent.critic_local.parameters(), 
                                             config.clip_grad_crit) # Clip Gradients
          agent.critic_optimizer.step()
          
          ##########################
          
          # ---------------------------- update actor ---------------------------- #
          # Compute actor loss
          all_actions_pred = [ self.maddpg_agent[i].actor_local(ob) if i == agent_number \
                     else self.maddpg_agent[i].actor_local(ob).detach()
                     for i, ob in enumerate(states) ]
          
          all_actions_pred = torch.cat(all_actions_pred, dim=1)
          actor_loss = -agent.critic_local(all_states, all_actions_pred).mean()
  
          # Minimize the loss
          agent.actor_optimizer.zero_grad()
          actor_loss.backward()
          if config.clip_grad_act:
              torch.nn.utils.clip_grad_norm_(agent.actor_local.parameters(), 
                                             config.clip_grad_act) # Clip Gradients
          agent.actor_optimizer.step()
          
          #####################
          
          self.soft_update(agent.critic_local, agent.critic_target, config.tau)
          self.soft_update(agent.actor_local, agent.actor_target, config.tau) 
  
      def soft_update(self, local_model, target_model, tau):
          """Soft update model parameters.
          θ_target = τ*θ_local + (1 - τ)*θ_target
  
          Params
          ======
              local_model: PyTorch model (weights will be copied from).
              target_model: PyTorch model (weights will be copied to).
              tau (float): interpolation parameter.
          """
          for target_param, local_param in zip(target_model.parameters(), local_model.parameters()):
              target_param.data.copy_(tau*local_param.data + (1.0-tau)*target_param.data)
              
      def save_agents(self):
          """Save Weights of every agent"""
          for idx, ddpg_agent in enumerate(self.maddpg_agent):
              torch.save(ddpg_agent.actor_local.state_dict(), f'MADDPG_chkpnt_actor_{idx}.pth')
              torch.save(ddpg_agent.critic_local.state_dict(),f'MADDPG_chkpnt_critic_{idx}.pth')
              torch.save(ddpg_agent.actor_target.state_dict(), f'MADDPG_chkpnt_actor_target_{idx}.pth')
              torch.save(ddpg_agent.critic_target.state_dict(),f'MADDPG_chkpnt_critic_target_{idx}.pth')
  
      def load_agents(self):
          """Load Weights of every agent"""
          for idx, ddpg_agent in enumerate(self.maddpg_agent):
              ddpg_agent.actor_local.load_state_dict(torch.load(f'MADDPG_chkpnt_actor_{idx}.pth'))
              ddpg_agent.critic_local.load_state_dict(torch.load(f'MADDPG_chkpnt_critic_{idx}.pth'))
              ddpg_agent.actor_target.load_state_dict(torch.load(f'MADDPG_chkpnt_actor_target_{idx}.pth'))
              ddpg_agent.critic_target.load_state_dict(torch.load(f'MADDPG_chkpnt_critic_target_{idx}.pth'))    