<filename>agent.py
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'))    
<filename>agent.py
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'))    
Deep Deterministic Policy Gradients (DDPG) Actor-Critic Agent. 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. # Actor Network (w/ Target Network) # Critic Network (w/ Target Network) # make sure local and target networks start with same weights # Noise process Returns actions for given state as per current policy. Returns actions for given state as per target policy. 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) Multi-Agent Deep Deterministic Policy Gradients (MADDPG) Actor-Critic Agent. 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. # Replay memory # Initialize time step (for updating every UPDATE_EVERY steps) Numpy -> Pytorch tensor Reset noise of all agents. Get noise of all agents. Get actions from all agents in the MADDPG object. Get target network actions from all the agents in the MADDPG object. Save experience in replay memory, and use random sample from buffer to learn. # Save experience / reward # Learn, if enough samples are available in memory 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). # ---------------------------- update critic ---------------------------- # # Get predicted next-state actions and Q values from target models # Compute Q targets for current states (y_i) # Compute critic loss # Minimize the loss # Clip Gradients ########################## # ---------------------------- update actor ---------------------------- # # Compute actor loss # Minimize the loss # Clip Gradients ##################### 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. Save Weights of every agent Load Weights of every agent