from __future__ import  print_function, division
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import torch.optim as optim
import copy

class CDC:

    def __init__(self, 
                actor,
                critic,
                p_lr = None,
                q_lr = None,
                gamma =0.99,
                ptau = 0.005,
                batch_size = 100,
                max_action = None,
                num_samples = 15,
                eta_coef = 1.0,
                action_dims = None,
                nu = 1.0,
                max_timesteps = 1e6,
                lambda_coef = 1.0,
                device = 'cpu',
                ):
        '''
            actor:  actor network
            critic: critic network
            p_lr:   policy learning rate
            q_lr:   Q learning rate
            gamma: reward discounting parameter
            ptau:  Interpolation factor in polyak averaging
            num_samples: number of actions samples
        '''
        # networks
        self.actor = actor
        self.critic = critic

        #####
        # params
        #####
        self.gamma = gamma
        self.ptau = ptau
        self.max_action = max_action
        self.batch_size = batch_size
        self.device = device
        self.nu = nu
        self.num_samples = num_samples
        self.eta_coef = eta_coef
        self.action_dims = np.prod(action_dims)
        self.total_grad_steps  = 0
        self.lambda_coef = lambda_coef

        #####
        # load target models.
        #####
        self.critic_target  = copy.deepcopy(self.critic)
        self.number_qs = self.critic.number_of_qs

        # Freeze target networks with respect to optimizers (only update via polyak averaging)
        for p in self.critic_target.parameters():
            p.requires_grad = False

        #####
        # Optims
        #####
        self.actor_optimizer  = optim.Adam(self.actor.parameters(), lr=p_lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=q_lr)

        print('-----------------------------')
        print('Optim Params')
        print("Actor:\n ",  self.actor_optimizer)
        print("Critic:\n ", self.critic_optimizer )
        print('-----------------------------')
        print('Hyper-parameters and flags details')
        print('batch_size: ', self.batch_size)
        print('nu: ', self.nu)
        print('action_dims: ', self.action_dims)
        print('num_samples: ', self.num_samples)
        print('number_qs: ', self.number_qs)
        print('lambda_coef: ', self.lambda_coef)
        print('eta_coef: ', self.eta_coef)
        print('-----------------------------')

    def select_action(self, obs, deterministic=False):
        '''
            return action
        '''
        # since there is a dropout in critic, make sure it is off during eval whenever drop > 0
        self.critic.eval()

        with torch.no_grad():
            # obs [D] ==> [1, D] ==> [num_samples, D]
            state = torch.FloatTensor(obs.reshape(1, -1)).repeat(self.num_samples, 1).to(self.device)
            # action [num_samples, action_dim]
            action = self.actor(state)
            # qis [num_qs, B, 1]
            qis  = self.critic(state, action)

            # qis [num_qs, B, 1] ==> q [B, 1]
            q = self.nu * qis.min(0)[0] + (1. - self.nu) * qis.max(0)[0]

            # now find max value.
            ind  = q.argmax(0)

        return action[ind].cpu().data.numpy().flatten()

    def lower_bound_Q_max(self, obs, next_obs, cQs_current=None, ct_loss=None):
        '''
            This function return loss between Q_pi(s, a~pi) without/or/and Q_b(s, a~b)
            cQs_current: [num_qs, B, 1]
            obs: [B, D]
            next_obs: [B, D]
        '''
        bsize = obs.shape[0]
        a_dim = self.action_dims

        # obs_rep [B, D] ==> view  [B * num_samples, D]
        # obs_rep = torch.repeat_interleave(obs, self.num_samples, 0)
        obs_rep = obs.unsqueeze(1).repeat(1, self.num_samples, 1).view(bsize * self.num_samples, -1)
        with torch.no_grad():
            # get actions from current policy (pi)
            curr_actions = self.actor(obs_rep)

        # curr_actions [B * num_samples, action_dim]
        # obs_rep [B * num_samples, obs_dim]
        # qXs:  [number_qs, B * num_samples, 1]
        # cQs_current: [num_qs, B, 1]
        qXs = self.critic(obs_rep, curr_actions)
        # qXs:  [number_qs, B * num_samples, 1] ==> [number_qs, B , num_samples] ==> max [number_qs, B] ==>  [number_qs, B, 1]
        max_qXs = qXs.reshape(self.number_qs, self.batch_size, self.num_samples).max(2)[0].reshape(self.number_qs, self.batch_size, 1)

        all_losses = []
        for k in range(self.number_qs):
            for i in range(self.number_qs):
                tm_qs = (torch.relu(max_qXs[i] - cQs_current[k])**2).mean()
                all_losses.append(tm_qs)

        mloss = sum(all_losses)/len(all_losses)

        return mloss

    def compute_critic_loss(self, obs, next_obs, action, reward, mask):
        '''
            Compute loss Q:
            t =  r + gamma * mask *( min_{1,2} Q_target(s', a') - alpha * log pi(s', a'))
        '''
        with torch.no_grad():

                ########
                # Target update
                ########
                # next_action [B, num_actions]
                bsize = next_obs.shape[0]
                next_obs_rep = torch.repeat_interleave(next_obs, self.num_samples, 0)

                # get next actions using current policy
                next_action_rep = self.actor(next_obs_rep)
                every_target_Qs = self.critic_target(next_obs_rep, next_action_rep)

                # we want to do max_a (min_i Q_i(a))
                # every_target_Qs [number_qs, B * samples , 1] ==> [B * samples , 1]
                min_target_qs = self.nu * every_target_Qs.min(0)[0] + (1. - self.nu) * every_target_Qs.max(0)[0]

                # [B * samples , 1] ==> reshape [B, samples] ==> max [B] ==> reshape [B, 1]
                target_Q = min_target_qs.reshape(self.batch_size, -1).max(1)[0].reshape(-1, 1)

                # backup: [B, 1]
                backup = reward + self.gamma  * mask * target_Q

        # 2.  Get current Q estimates
        cQs = self.critic(obs, action)

        # 3. Compute critic loss
        # even we picked min Q, we still need to backprob to both Qs
        all_critic_losses = []
        for ci in cQs:
            all_critic_losses.append(F.mse_loss(ci, backup))

        critic_loss =  sum(all_critic_losses)

        return critic_loss, cQs

    def compute_loss_pi(self, obs, actions):
        '''
            Compute pi loss
               loss = alpha * log_pi - min_{1,2} Q(s, pi)
        '''
        pi_action, lopprobs = self.actor(obs, gt_actions=actions, with_log_mle=True)
        qs_pi = self.critic(obs, pi_action)
        # q_pi  = qs_pi.min(0)[0]
        q_pi  = self.nu * qs_pi.min(0)[0] + (1. - self.nu) * qs_pi.max(0)[0]
        loss_pi = (- q_pi -  self.lambda_coef * lopprobs ).mean()

        return loss_pi, lopprobs

    def train(self, replay_buffer=None,
                    iterations=None,
                    burn_in_passed=False):

        '''
            inputs:
                replay_buffer
            outputs:
                 
        '''
        actor_loss_out = 0.0
        critic_loss_out = 0.0
        qs_loss_out = 0
        lgprob_out = 0

        # since there is a dropout in critic and
        # make sure it is on whenever drop > 0
        self.critic.train()

        for it in range(iterations):
            ########
            # Sample replay buffer
            ########
            x, y, u, r, d = replay_buffer.sample(self.batch_size)
            obs = torch.FloatTensor(x).to(self.device)
            next_obs = torch.FloatTensor(y).to(self.device)
            action = torch.FloatTensor(u).to(self.device)
            reward = torch.FloatTensor(r).to(self.device)
            mask = torch.FloatTensor(1. - d).to(self.device)

            ########
            # policy updates
            ########
            # Freeze Q-networks so you don't waste computational effort
            # computing gradients for them during the policy learning step.
            for p in self.critic.parameters():
                p.requires_grad = False

            actor_loss, lgprp = self.compute_loss_pi(obs, action)

            actor_loss_out += actor_loss.item()
            lgprob_out += lgprp.mean().item()

            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            # Unfreeze Q-networks so you can optimize it at next DDPG step.
            for p in self.critic.parameters():
                p.requires_grad = True

            ########
            # critic updates
            ########
            # 1. Compute critic loss
            # even we picked min Q, we still need to backprob to both Qs
            critic_loss, ccQs = self.compute_critic_loss(obs, next_obs, action, reward, mask)
            critic_loss_out += critic_loss.item()

            # 2. now get difference between Qs
            qs_loss = self.lower_bound_Q_max(obs, next_obs, cQs_current=ccQs, ct_loss=critic_loss_out)
            qs_loss_out += qs_loss.item()

            crt_losses =  critic_loss + self.eta_coef * qs_loss

            # 3. Optimize the critic
            self.critic_optimizer.zero_grad()
            crt_losses.backward()
            self.critic_optimizer.step()

            ########
            # target updates
            ########
            with torch.no_grad():

                # Update the frozen target models
                for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
                    # Use an in-place operations "mul_", "add_" to update target
                    # params, as opposed to "mul" and "add", which would make new tensors.
                    target_param.data.mul_((1 - self.ptau))
                    target_param.data.add_(self.ptau * param.data)

            self.total_grad_steps += 1

        out = {}
        out['lg_loss']     = -lgprob_out/iterations
        out['critic_loss'] = critic_loss_out/iterations
        out['actor_loss']  = actor_loss_out/iterations
        out['qs_loss']     = qs_loss_out/iterations

        return out
from __future__ import  print_function, division
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import torch.optim as optim
import copy

class CDC:

    def __init__(self, 
                actor,
                critic,
                p_lr = None,
                q_lr = None,
                gamma =0.99,
                ptau = 0.005,
                batch_size = 100,
                max_action = None,
                num_samples = 15,
                eta_coef = 1.0,
                action_dims = None,
                nu = 1.0,
                max_timesteps = 1e6,
                lambda_coef = 1.0,
                device = 'cpu',
                ):
        '''
            actor:  actor network
            critic: critic network
            p_lr:   policy learning rate
            q_lr:   Q learning rate
            gamma: reward discounting parameter
            ptau:  Interpolation factor in polyak averaging
            num_samples: number of actions samples
        '''
        # networks
        self.actor = actor
        self.critic = critic

        #####
        # params
        #####
        self.gamma = gamma
        self.ptau = ptau
        self.max_action = max_action
        self.batch_size = batch_size
        self.device = device
        self.nu = nu
        self.num_samples = num_samples
        self.eta_coef = eta_coef
        self.action_dims = np.prod(action_dims)
        self.total_grad_steps  = 0
        self.lambda_coef = lambda_coef

        #####
        # load target models.
        #####
        self.critic_target  = copy.deepcopy(self.critic)
        self.number_qs = self.critic.number_of_qs

        # Freeze target networks with respect to optimizers (only update via polyak averaging)
        for p in self.critic_target.parameters():
            p.requires_grad = False

        #####
        # Optims
        #####
        self.actor_optimizer  = optim.Adam(self.actor.parameters(), lr=p_lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=q_lr)

        print('-----------------------------')
        print('Optim Params')
        print("Actor:\n ",  self.actor_optimizer)
        print("Critic:\n ", self.critic_optimizer )
        print('-----------------------------')
        print('Hyper-parameters and flags details')
        print('batch_size: ', self.batch_size)
        print('nu: ', self.nu)
        print('action_dims: ', self.action_dims)
        print('num_samples: ', self.num_samples)
        print('number_qs: ', self.number_qs)
        print('lambda_coef: ', self.lambda_coef)
        print('eta_coef: ', self.eta_coef)
        print('-----------------------------')

    def select_action(self, obs, deterministic=False):
        '''
            return action
        '''
        # since there is a dropout in critic, make sure it is off during eval whenever drop > 0
        self.critic.eval()

        with torch.no_grad():
            # obs [D] ==> [1, D] ==> [num_samples, D]
            state = torch.FloatTensor(obs.reshape(1, -1)).repeat(self.num_samples, 1).to(self.device)
            # action [num_samples, action_dim]
            action = self.actor(state)
            # qis [num_qs, B, 1]
            qis  = self.critic(state, action)

            # qis [num_qs, B, 1] ==> q [B, 1]
            q = self.nu * qis.min(0)[0] + (1. - self.nu) * qis.max(0)[0]

            # now find max value.
            ind  = q.argmax(0)

        return action[ind].cpu().data.numpy().flatten()

    def lower_bound_Q_max(self, obs, next_obs, cQs_current=None, ct_loss=None):
        '''
            This function return loss between Q_pi(s, a~pi) without/or/and Q_b(s, a~b)
            cQs_current: [num_qs, B, 1]
            obs: [B, D]
            next_obs: [B, D]
        '''
        bsize = obs.shape[0]
        a_dim = self.action_dims

        # obs_rep [B, D] ==> view  [B * num_samples, D]
        # obs_rep = torch.repeat_interleave(obs, self.num_samples, 0)
        obs_rep = obs.unsqueeze(1).repeat(1, self.num_samples, 1).view(bsize * self.num_samples, -1)
        with torch.no_grad():
            # get actions from current policy (pi)
            curr_actions = self.actor(obs_rep)

        # curr_actions [B * num_samples, action_dim]
        # obs_rep [B * num_samples, obs_dim]
        # qXs:  [number_qs, B * num_samples, 1]
        # cQs_current: [num_qs, B, 1]
        qXs = self.critic(obs_rep, curr_actions)
        # qXs:  [number_qs, B * num_samples, 1] ==> [number_qs, B , num_samples] ==> max [number_qs, B] ==>  [number_qs, B, 1]
        max_qXs = qXs.reshape(self.number_qs, self.batch_size, self.num_samples).max(2)[0].reshape(self.number_qs, self.batch_size, 1)

        all_losses = []
        for k in range(self.number_qs):
            for i in range(self.number_qs):
                tm_qs = (torch.relu(max_qXs[i] - cQs_current[k])**2).mean()
                all_losses.append(tm_qs)

        mloss = sum(all_losses)/len(all_losses)

        return mloss

    def compute_critic_loss(self, obs, next_obs, action, reward, mask):
        '''
            Compute loss Q:
            t =  r + gamma * mask *( min_{1,2} Q_target(s', a') - alpha * log pi(s', a'))
        '''
        with torch.no_grad():

                ########
                # Target update
                ########
                # next_action [B, num_actions]
                bsize = next_obs.shape[0]
                next_obs_rep = torch.repeat_interleave(next_obs, self.num_samples, 0)

                # get next actions using current policy
                next_action_rep = self.actor(next_obs_rep)
                every_target_Qs = self.critic_target(next_obs_rep, next_action_rep)

                # we want to do max_a (min_i Q_i(a))
                # every_target_Qs [number_qs, B * samples , 1] ==> [B * samples , 1]
                min_target_qs = self.nu * every_target_Qs.min(0)[0] + (1. - self.nu) * every_target_Qs.max(0)[0]

                # [B * samples , 1] ==> reshape [B, samples] ==> max [B] ==> reshape [B, 1]
                target_Q = min_target_qs.reshape(self.batch_size, -1).max(1)[0].reshape(-1, 1)

                # backup: [B, 1]
                backup = reward + self.gamma  * mask * target_Q

        # 2.  Get current Q estimates
        cQs = self.critic(obs, action)

        # 3. Compute critic loss
        # even we picked min Q, we still need to backprob to both Qs
        all_critic_losses = []
        for ci in cQs:
            all_critic_losses.append(F.mse_loss(ci, backup))

        critic_loss =  sum(all_critic_losses)

        return critic_loss, cQs

    def compute_loss_pi(self, obs, actions):
        '''
            Compute pi loss
               loss = alpha * log_pi - min_{1,2} Q(s, pi)
        '''
        pi_action, lopprobs = self.actor(obs, gt_actions=actions, with_log_mle=True)
        qs_pi = self.critic(obs, pi_action)
        # q_pi  = qs_pi.min(0)[0]
        q_pi  = self.nu * qs_pi.min(0)[0] + (1. - self.nu) * qs_pi.max(0)[0]
        loss_pi = (- q_pi -  self.lambda_coef * lopprobs ).mean()

        return loss_pi, lopprobs

    def train(self, replay_buffer=None,
                    iterations=None,
                    burn_in_passed=False):

        '''
            inputs:
                replay_buffer
            outputs:
                 
        '''
        actor_loss_out = 0.0
        critic_loss_out = 0.0
        qs_loss_out = 0
        lgprob_out = 0

        # since there is a dropout in critic and
        # make sure it is on whenever drop > 0
        self.critic.train()

        for it in range(iterations):
            ########
            # Sample replay buffer
            ########
            x, y, u, r, d = replay_buffer.sample(self.batch_size)
            obs = torch.FloatTensor(x).to(self.device)
            next_obs = torch.FloatTensor(y).to(self.device)
            action = torch.FloatTensor(u).to(self.device)
            reward = torch.FloatTensor(r).to(self.device)
            mask = torch.FloatTensor(1. - d).to(self.device)

            ########
            # policy updates
            ########
            # Freeze Q-networks so you don't waste computational effort
            # computing gradients for them during the policy learning step.
            for p in self.critic.parameters():
                p.requires_grad = False

            actor_loss, lgprp = self.compute_loss_pi(obs, action)

            actor_loss_out += actor_loss.item()
            lgprob_out += lgprp.mean().item()

            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            self.actor_optimizer.step()

            # Unfreeze Q-networks so you can optimize it at next DDPG step.
            for p in self.critic.parameters():
                p.requires_grad = True

            ########
            # critic updates
            ########
            # 1. Compute critic loss
            # even we picked min Q, we still need to backprob to both Qs
            critic_loss, ccQs = self.compute_critic_loss(obs, next_obs, action, reward, mask)
            critic_loss_out += critic_loss.item()

            # 2. now get difference between Qs
            qs_loss = self.lower_bound_Q_max(obs, next_obs, cQs_current=ccQs, ct_loss=critic_loss_out)
            qs_loss_out += qs_loss.item()

            crt_losses =  critic_loss + self.eta_coef * qs_loss

            # 3. Optimize the critic
            self.critic_optimizer.zero_grad()
            crt_losses.backward()
            self.critic_optimizer.step()

            ########
            # target updates
            ########
            with torch.no_grad():

                # Update the frozen target models
                for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
                    # Use an in-place operations "mul_", "add_" to update target
                    # params, as opposed to "mul" and "add", which would make new tensors.
                    target_param.data.mul_((1 - self.ptau))
                    target_param.data.add_(self.ptau * param.data)

            self.total_grad_steps += 1

        out = {}
        out['lg_loss']     = -lgprob_out/iterations
        out['critic_loss'] = critic_loss_out/iterations
        out['actor_loss']  = actor_loss_out/iterations
        out['qs_loss']     = qs_loss_out/iterations

        return out
actor:  actor network             critic: critic network             p_lr:   policy learning rate             q_lr:   Q learning rate             gamma: reward discounting parameter             ptau:  Interpolation factor in polyak averaging             num_samples: number of actions samples # networks ##### # params ##### ##### # load target models. ##### # Freeze target networks with respect to optimizers (only update via polyak averaging) ##### # Optims ##### return action # since there is a dropout in critic, make sure it is off during eval whenever drop > 0 # obs [D] ==> [1, D] ==> [num_samples, D] # action [num_samples, action_dim] # qis [num_qs, B, 1] # qis [num_qs, B, 1] ==> q [B, 1] # now find max value. This function return loss between Q_pi(s, a~pi) without/or/and Q_b(s, a~b)             cQs_current: [num_qs, B, 1]             obs: [B, D]             next_obs: [B, D] # obs_rep [B, D] ==> view  [B * num_samples, D] # obs_rep = torch.repeat_interleave(obs, self.num_samples, 0) # get actions from current policy (pi) # curr_actions [B * num_samples, action_dim] # obs_rep [B * num_samples, obs_dim] # qXs:  [number_qs, B * num_samples, 1] # cQs_current: [num_qs, B, 1] # qXs:  [number_qs, B * num_samples, 1] ==> [number_qs, B , num_samples] ==> max [number_qs, B] ==>  [number_qs, B, 1] Compute loss Q:             t =  r + gamma * mask *( min_{1,2} Q_target(s', a') - alpha * log pi(s', a')) ######## # Target update ######## # next_action [B, num_actions] # get next actions using current policy # we want to do max_a (min_i Q_i(a)) # every_target_Qs [number_qs, B * samples , 1] ==> [B * samples , 1] # [B * samples , 1] ==> reshape [B, samples] ==> max [B] ==> reshape [B, 1] # backup: [B, 1] # 2.  Get current Q estimates # 3. Compute critic loss # even we picked min Q, we still need to backprob to both Qs Compute pi loss                loss = alpha * log_pi - min_{1,2} Q(s, pi) # q_pi  = qs_pi.min(0)[0] inputs:                 replay_buffer             outputs: # since there is a dropout in critic and # make sure it is on whenever drop > 0 ######## # Sample replay buffer ######## ######## # policy updates ######## # Freeze Q-networks so you don't waste computational effort # computing gradients for them during the policy learning step. # Unfreeze Q-networks so you can optimize it at next DDPG step. ######## # critic updates ######## # 1. Compute critic loss # even we picked min Q, we still need to backprob to both Qs # 2. now get difference between Qs # 3. Optimize the critic ######## # target updates ######## # Update the frozen target models # Use an in-place operations "mul_", "add_" to update target # params, as opposed to "mul" and "add", which would make new tensors.