<file_sep>import torch as K
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np
import copy 

from olbr.agents.basic import BackwardDyn

import pdb

def soft_update(target, source, tau):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)

def hard_update(target, source):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(param.data)


class PPO_BD(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, 
                 clip_param, ppo_epoch, num_mini_batch, value_loss_coef, entropy_coef,
                 eps=None, max_grad_norm=None, use_clipped_value_loss=True,
                 out_func=K.sigmoid, discrete=True, agent_id=0, object_Qfunc=None, backward_dyn=None, 
                 object_policy=None, reward_fun=None, masked_with_r=False, dtype=K.float32, device="cuda",
                 ):

        super(PPO_BD, self).__init__()

        self.observation_space = observation_space
        self.action_space = action_space
        
        optimizer, lr = optimizer
        actor_lr, _ = lr

        self.clip_param = clip_param
        self.ppo_epoch = ppo_epoch
        self.num_mini_batch = num_mini_batch

        self.value_loss_coef = value_loss_coef
        self.entropy_coef = entropy_coef

        self.max_grad_norm = max_grad_norm
        self.use_clipped_value_loss = use_clipped_value_loss

        self.discrete = discrete
        self.agent_id = agent_id

        self.object_Qfunc = object_Qfunc
        self.object_policy = object_policy

        self.dtype = dtype
        self.device = device
        self.masked_with_r = masked_with_r
        self.gamma = 0.98
        self.tau = 0.05
        self.loss_func = F.mse_loss

        # model initialization
        self.entities = []
        
        # actors and critics
        self.actors = []
        self.critics = []
        self.optims = []

        self.actors.append(Actor(observation_space, action_space[agent_id], discrete, out_func).to(device))
        self.critics.append(Critic(observation_space, None).to(device))

        self.optims.append(optimizer(self.critics[0].parameters(), lr=actor_lr, eps=eps))
        self.optims[0].add_param_group({"params": self.actors[0].parameters()})
        
        self.entities.extend(self.actors)
        self.entities.extend(self.critics)
        self.entities.extend(self.optims) 

        # backward dynamics model
        if backward_dyn is not None:
            self.backward = backward_dyn
            self.backward_optim = optimizer(self.backward.parameters(), lr = actor_lr)
            self.entities.append(self.backward)
            self.entities.append(self.backward_optim)

        # Learnt Q function for object
        if self.object_Qfunc is not None:
            self.object_Qfunc_target = copy.deepcopy(self.object_Qfunc)
            self.object_Qfunc_optim = optimizer(self.object_Qfunc.parameters(), lr = actor_lr)
            self.entities.append(self.object_Qfunc)
            self.entities.append(self.object_Qfunc_target)
            self.entities.append(self.object_Qfunc_optim)

        # Learnt policy for object
        if self.object_policy is not None:
            self.object_policy_target = copy.deepcopy(self.object_policy)
            self.object_policy_optim = optimizer(self.object_policy.parameters(), lr = actor_lr)
            self.entities.append(self.object_policy)
            self.entities.append(self.object_policy_target)
            self.entities.append(self.object_policy_optim)

        if reward_fun is not None:
            self.get_obj_reward = reward_fun
        else:
            self.get_obj_reward = self.reward_fun
        
        print('clipped between -1 and 0, and masked with abs(r), and + r')

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.optims[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.optims[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False):
        with K.no_grad():
            value = self.critics[0](state.to(self.device))
            action_dist = self.actors[0](state.to(self.device))

        if exploration:
            action = action_dist.sample()
        else:
            action = action_dist.mode()

        #action = action.clamp(int(self.action_space[self.agent_id].low[0]), int(self.action_space[self.agent_id].high[0]))

        return value, action, action_dist.log_probs(action)

    def get_obj_action(self, state, exploration=False):
        self.object_policy.eval()
        with K.no_grad():
            mu = self.object_policy(state.to(self.device))
        self.object_policy.train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[1].low[0]), int(self.action_space[1].high[0]))

        return mu

    def evaluate_actions(self, state, action, get_preactivations=False):       
        value = self.critics[0](state.to(self.device))
        action_dist = self.actors[0](state.to(self.device))

        action_log_probs = action_dist.log_probs(action)
        dist_entropy = action_dist.entropies().mean()

        if get_preactivations:
            action_preactivations = self.actors[0].get_preactivations(state.to(self.device))
        else:
            action_preactivations = None

        return value, action_log_probs, dist_entropy, action_preactivations

    def reward_fun(self, state, next_state):
        with K.no_grad():
            action = self.backward(state.to(self.device), next_state.to(self.device))
            opt_action = self.object_policy(state.to(self.device))

            reward = self.object_Qfunc(state.to(self.device), action) - self.object_Qfunc(state.to(self.device), opt_action)
        return reward.clamp(min=-1.0, max=0.0)

    def update(self, rollouts):
        advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

        value_loss_epoch = 0
        action_loss_epoch = 0
        dist_entropy_epoch = 0

        for e in range(self.ppo_epoch):
            data_generator = rollouts.feed_forward_generator(advantages, self.num_mini_batch, mini_batch_size=None)

            for sample in data_generator:
                obs_batch, actions_batch, value_preds_batch, return_batch, old_action_log_probs_batch, adv_targ = sample

                # Reshape to do in a single forward pass for all steps
                values, action_log_probs, dist_entropy, action_preactivations  = self.evaluate_actions(obs_batch, actions_batch, False)

                ratio = K.exp(action_log_probs - old_action_log_probs_batch)
                surr1 = ratio * adv_targ
                surr2 = K.clamp(ratio, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ
                action_loss = -K.min(surr1, surr2).mean()
                #action_loss += (action_preactivations**2).mean()*0.001

                if self.use_clipped_value_loss:
                    value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param, self.clip_param)
                    value_losses = (values - return_batch).pow(2)
                    value_losses_clipped = (value_pred_clipped - return_batch).pow(2)
                    value_loss = 0.5 * K.max(value_losses, value_losses_clipped).mean()
                else:
                    value_loss = 0.5 * (return_batch - values).pow(2).mean()

                self.optims[0].zero_grad()
                (value_loss * self.value_loss_coef + action_loss - dist_entropy * self.entropy_coef).backward()
                nn.utils.clip_grad_norm_(self.critics[0].parameters(), self.max_grad_norm)
                nn.utils.clip_grad_norm_(self.actors[0].parameters(), self.max_grad_norm)
                self.optims[0].step()

                value_loss_epoch += value_loss.item()
                action_loss_epoch += action_loss.item()
                dist_entropy_epoch += dist_entropy.item()

        num_updates = self.ppo_epoch * self.num_mini_batch

        value_loss_epoch /= num_updates
        action_loss_epoch /= num_updates
        dist_entropy_epoch /= num_updates

        return value_loss_epoch, action_loss_epoch, dist_entropy_epoch

    def update_backward(self, batch, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]
        
        s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        if normalizer[1] is not None:
            s2 = normalizer[1].preprocess(s2)
            s2_ = normalizer[1].preprocess(s2_)

        a2_pred = self.backward(s2, s2_)

        loss_backward = self.loss_func(a2_pred, a2)

        self.backward_optim.zero_grad()
        loss_backward.backward()
        self.backward_optim.step()

        return loss_backward.item()

    def update_object_parameters(self, batch, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)
        
        s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        if normalizer[1] is not None:
            s2 = normalizer[1].preprocess(s2)
            s2_ = normalizer[1].preprocess(s2_)

        s, s_, a = s2, s2_, a2
        a_ = self.object_policy_target(s_)
    
        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device).unsqueeze(1)

        Q = self.object_Qfunc(s, a)       
        V = self.object_Qfunc_target(s_, a_).detach()

        target_Q = (V * self.gamma) + r
        target_Q = target_Q.clamp(-1./(1.-self.gamma), 0.)

        loss_critic = self.loss_func(Q, target_Q)

        self.object_Qfunc_optim.zero_grad()
        loss_critic.backward()
        self.object_Qfunc_optim.step()

        a = self.object_policy(s)

        loss_actor = -self.object_Qfunc(s, a).mean()
        loss_actor += (self.object_policy(s)**2).mean()*1

        self.object_policy_optim.zero_grad()        
        loss_actor.backward()
        self.object_policy_optim.step()
                
        return loss_critic.item(), loss_actor.item()

    def update_object_target(self):

        soft_update(self.object_policy_target, self.object_policy, self.tau)
        soft_update(self.object_Qfunc_target, self.object_Qfunc, self.tau)
<file_sep>import numpy as np
import scipy as sc
import time
import imageio

import torch as K
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import gym_wmgds as gym

from olbr.algorithms import DDPG
from olbr.experience import ReplayMemory, Transition, Normalizer
from olbr.exploration import Noise
from olbr.utils import Saver, Summarizer, get_noise_scale, get_params, running_mean
from olbr.agents.basic import Actor 
from olbr.agents.basic import Critic

import pdb

import matplotlib
import matplotlib.pyplot as plt

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

def init(config):
    
    if config['resume'] != '':
        resume_path = config['resume']
        saver = Saver(config)
        config, start_episode, save_dict = saver.resume_ckpt()
        config['resume'] = resume_path
    else:
        start_episode = 0
        
    #hyperparameters
    ENV_NAME = config['env_id'] #'simple_spread'
    SEED = config['random_seed'] # 1

    GAMMA = config['gamma'] # 0.95
    TAU = config['tau'] # 0.01

    ACTOR_LR = config['plcy_lr'] # 0.01
    CRITIC_LR = config['crtc_lr'] # 0.01

    MEM_SIZE = config['buffer_length'] # 1000000 

    REGULARIZATION = config['regularization'] # True
    NORMALIZED_REWARDS = config['reward_normalization'] # True

    env = gym.make(ENV_NAME)
    env.seed(SEED)

    observation_space = env.observation_space.spaces['observation'].shape[0] + env.observation_space.spaces['desired_goal'].shape[0]
    action_space = env.action_space
    if env.action_space.low[0] == -1 and env.action_space.high[0] == 1:
        OUT_FUNC = K.tanh 
    elif env.action_space.low[0] == 0 and env.action_space.high[0] == 1:
        OUT_FUNC = K.sigmoid
    else:
        OUT_FUNC = K.sigmoid

    K.manual_seed(SEED)
    np.random.seed(SEED)
    
    MODEL = DDPG

    if config['verbose'] > 1:
        # utils
        summaries = (Summarizer(config['dir_summary_train'], config['port'], config['resume']),
                    Summarizer(config['dir_summary_test'], config['port'], config['resume']))
        saver = Saver(config)
    else:
        summaries = None
        saver = None

    #exploration initialization
    noise = Noise(action_space.shape[0], sigma=0.2, eps=0.3)
    #noise = OUNoise(action_space.shape[0])

    #model initialization
    optimizer = (optim.Adam, (ACTOR_LR, CRITIC_LR)) # optimiser func, (actor_lr, critic_lr)
    loss_func = F.mse_loss
    model = MODEL(observation_space, action_space, optimizer, 
                  Actor, Critic, loss_func, GAMMA, TAU, out_func=OUT_FUNC,
                  discrete=False, regularization=REGULARIZATION, normalized_rewards=NORMALIZED_REWARDS)
    
    if config['resume'] != '':
        for i, param in enumerate(save_dict['model_params']):
            model.entities[i].load_state_dict(param)
    
    #memory initilization
    memory = ReplayMemory(MEM_SIZE)

    normalizer = (Normalizer(), None)

    experiment_args = (env, memory, noise, config, summaries, saver, start_episode, normalizer)
          
    return model, experiment_args

def rollout(env, model, noise, normalizer=None, render=False, nb_objects=1):
    trajectory = []

    # monitoring variables
    episode_reward = 0
    frames = []
    
    env.env.nb_objects = nb_objects
    state = env.reset()
    achieved_init = state['achieved_goal']
    
    for i_step in range(env._max_episode_steps):

        model.to_cpu()

        obs = K.tensor(state['observation'][0], dtype=K.float32).unsqueeze(0)
        goal = K.tensor(state['desired_goal'], dtype=K.float32).unsqueeze(0)

        obs_goal = K.cat([obs, goal], dim=-1)
        # Observation normalization
        if normalizer[0] is not None:
            obs_goal = normalizer[0].preprocess_with_update(obs_goal)
            #_ = normalizer.preprocess_with_update(obs_goal)

        action = model.select_action(obs_goal, noise)

        next_state, reward, done, info = env.step(action.squeeze(0).numpy())
        reward = K.tensor(reward, dtype=dtype).view(1,1)
        if normalizer[1] is not None:
            reward = normalizer[1].preprocess_with_update(reward)

        # for monitoring
        episode_reward += reward
        trajectory.append((state.copy(), action, reward, next_state.copy(), done))

        # Move to the next state
        state = next_state

        # Record frames
        if render:
            frames.append(env.render(mode='rgb_array')[0])
        
        #achieved_final = state['achieved_goal']
        #moved_step[np.where(np.logical_and(np.linalg.norm((achieved_init - achieved_final), axis=1) > 1e-4, moved_step==-1))[0]] = i_step

    achieved_final = state['achieved_goal']
    moved_index = np.concatenate((K.zeros(1), np.where((achieved_init[1::] != achieved_final[1::]).all(axis=1))[0]+1)).astype(int)
    #moved_index = np.where(np.linalg.norm((achieved_init - achieved_final), axis=1) > 1e-4)[0]

    if render:
        print(moved_index)


    return trajectory, episode_reward, info['is_success'], frames, moved_index

def run(model, experiment_args, train=True):

    total_time_start =  time.time()

    env, memory, noise, config, summaries, saver, start_episode, normalizer = experiment_args
    
    start_episode = start_episode if train else 0
    NUM_EPISODES = config['n_episodes'] if train else config['n_episodes_test'] 
    EPISODE_LENGTH = env._max_episode_steps #config['episode_length'] if train else config['episode_length_test'] 
    
    episode_reward_all = []
    episode_success_all = []
    critic_losses = []
    actor_losses = []
    
    max_nb_objects = config['max_nb_objects']
        
    for i_episode in range(start_episode, NUM_EPISODES):
        
        episode_time_start = time.time()
        #noise.scale = get_noise_scale(i_episode, config)
        if train:
            for i_cycle in range(50):
            
                trajectories = []
                moved_indices = []
                for i_rollout in range(1):
                    # Initialize the environment and state
                    nb_objects = np.random.randint(1, max_nb_objects+1)
                    trajectory, _, _, _, moved_index = rollout(env, model, noise, normalizer, render=(i_cycle%10==1000), nb_objects=nb_objects)
                    trajectories.append(trajectory)
                    moved_indices.append(moved_index)

                for trajectory, moved_index in zip(trajectories, moved_indices):

                    T = len(trajectory) - 1
                    t_samples = np.random.randint(T, size=(max_nb_objects,T))
                    future_offset = np.random.uniform(size=(max_nb_objects,T)) * (T - t_samples)
                    future_offset = future_offset.astype(int)
                    her_sample = (np.random.uniform(size=(max_nb_objects,T)) < 0.8)
                    future_t = (t_samples + 1 + future_offset)
                    
                    # for i_step in range(len(trajectory)):
                    #     state, action, reward, next_state, done = trajectory[i_step]
                        
                    #     obs = K.tensor(state['observation'][0], dtype=K.float32).unsqueeze(0)
                    #     goal = K.tensor(state['desired_goal'], dtype=K.float32).unsqueeze(0)

                    #     next_obs = K.tensor(next_state['observation'][0], dtype=K.float32).unsqueeze(0)
                    #     next_achieved = K.tensor(next_state['achieved_goal'][0], dtype=K.float32).unsqueeze(0)
                        
                    #     # regular sample
                    #     obs_goal = K.cat([obs, goal], dim=-1)
                    #     if done:
                    #         next_obs_goal = None
                    #     else:
                    #         next_obs_goal = K.cat([next_obs, goal], dim=-1)
                        
                    #     memory.push(obs_goal, action, next_obs_goal, reward)
                    # <-- end loop: i_step
                    
                    for i_object in moved_index:
                        for i_step in t_samples[i_object]:
                            state, action, reward, next_state, done = trajectory[i_step]

                            #pdb.set_trace()
                            obs = K.tensor(state['observation'][i_object], dtype=K.float32).unsqueeze(0)
                            goal = K.tensor(state['desired_goal'], dtype=K.float32).unsqueeze(0)
                                
                            next_obs = K.tensor(next_state['observation'][i_object], dtype=K.float32).unsqueeze(0)
                            next_achieved = K.tensor(next_state['achieved_goal'][i_object], dtype=K.float32).unsqueeze(0)

                            if her_sample[i_object, i_step]:
                                _, _, _, next_state, _ = trajectory[future_t[i_object, i_step]]
                                aux_goal = K.tensor(next_state['achieved_goal'][i_object], dtype=K.float32).unsqueeze(0)
                                obs_goal = K.cat([obs, aux_goal], dim=-1)

                                if done:
                                    next_obs_goal = None
                                else:
                                    next_obs_goal = K.cat([next_obs, aux_goal], dim=-1)
                                        
                                reward = env.compute_reward(next_achieved, aux_goal, None)
                                reward = K.tensor(reward, dtype=dtype).view(1,1)
                                if normalizer[1] is not None:
                                    reward = normalizer[1].preprocess_with_update(reward)
                            else:
                                # regular sample
                                obs_goal = K.cat([obs, goal], dim=-1)
                                if done:
                                    next_obs_goal = None
                                else:
                                    next_obs_goal = K.cat([next_obs, goal], dim=-1)
                                    
                            memory.push(obs_goal, action, next_obs_goal, reward)    
                        # <-- end loop: i_step
                    # <-- end loop: i_object  
                # <-- end loop: i_rollout 

                for i_batch in range(40):
                    if len(memory) > config['batch_size']-1:
                        
                        model.to_cuda()  
                        batch = Transition(*zip(*memory.sample(config['batch_size'])))
                        critic_loss, actor_loss = model.update_parameters(batch, normalizer)

                        if i_batch == 39:
                            critic_losses.append(critic_loss)
                            actor_losses.append(actor_loss)

                #print(normalizer.get_stats())

            # <-- end loop: i_cycle
        
        plot_durations(np.asarray(critic_losses), np.asarray(actor_losses))
        #pdb.set_trace()
        episode_reward_cycle = []
        episode_succeess_cycle = []
        for i_rollout in range(10):
            # Initialize the environment and state
            render = config['render'] > 0 and i_episode % config['render'] == 0
            _, episode_reward, success, frames, _ = rollout(env, model, noise=False, normalizer=normalizer, render=render, nb_objects=1)
                
            # Save gif
            dir_monitor = config['dir_monitor_train'] if train else config['dir_monitor_test']
            if config['render'] > 0 and i_episode % config['render'] == 0:
                imageio.mimsave('{}/{}.gif'.format(dir_monitor, i_episode), frames)

            episode_reward_cycle.append(episode_reward)
            episode_succeess_cycle.append(success)
        # <-- end loop: i_rollout 
            
        ### MONITORIRNG ###
        episode_reward_all.append(np.mean(episode_reward_cycle))
        episode_success_all.append(np.mean(episode_succeess_cycle))

        plot_durations(np.asarray(episode_reward_all), np.asarray(episode_success_all))
        
        if config['verbose'] > 0:
        # Printing out
            if (i_episode+1)%1 == 0:
                print("==> Episode {} of {}".format(i_episode + 1, NUM_EPISODES))
                print('  | Id exp: {}'.format(config['exp_id']))
                print('  | Exp description: {}'.format(config['exp_descr']))
                print('  | Env: {}'.format(config['env_id']))
                print('  | Process pid: {}'.format(config['process_pid']))
                print('  | Tensorboard port: {}'.format(config['port']))
                print('  | Episode total reward: {}'.format(episode_reward))
                print('  | Running mean of total reward: {}'.format(episode_reward_all[-1]))
                print('  | Success rate: {}'.format(episode_success_all[-1]))
                #print('  | Running mean of total reward: {}'.format(running_mean(episode_reward_all)[-1]))
                print('  | Time episode: {}'.format(time.time()-episode_time_start))
                print('  | Time total: {}'.format(time.time()-total_time_start))
                        
        if config['verbose'] > 0:    
            ep_save = i_episode+1 if (i_episode == NUM_EPISODES-1) else None  
            is_best_save = None
            is_best_avg_save = None   
                
            if (not train) or ((np.asarray([ep_save, is_best_save, is_best_avg_save]) == None).sum() == 3):
                to_save = False
            else:
                model.to_cpu()
                saver.save_checkpoint(save_dict   = {'model_params': [entity.state_dict() for entity in model.entities]},
                                      episode     = ep_save,
                                      is_best     = is_best_save,
                                      is_best_avg = is_best_avg_save
                                      )
                to_save = True
    
            #if (i_episode+1)%100 == 0:
            #    summary = summaries[0] if train else summaries[1]
            #    summary.update_log(i_episode, 
            #                    episode_reward, 
            #                    list(episode_reward.reshape(-1,)), 
            #                    critic_loss        = critic_loss, 
            #                    actor_loss         = actor_loss,
            #                    to_save            = to_save
            #                    )
            
    # <-- end loop: i_episode
    if train:
        print('Training completed')
    else:
        print('Test completed')
    
    return episode_reward_all

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display

plt.ion()

def plot_durations(p, r):
    plt.figure(2)
    plt.clf()
    plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Duration')
    plt.plot(p)
    plt.plot(r)

    plt.pause(0.001)  # pause a bit so that plots are updated
    if is_ipython:
        display.clear_output(wait=True)
        display.display(plt.gcf())

if __name__ == '__main__':

    monitor_macddpg_p2 = []
    monitor_macddpg_p2_test = []
    for i in range(0,5):
        config = get_params(args=['--exp_id','MACDDPG_P2_120K_'+ str(i+1), 
                                '--random_seed', str(i+1), 
                                '--agent_alg', 'MACDDPG',
                                '--protocol_type', str(2),
                                '--n_episodes', '120000',
                                '--verbose', '2',
                                ]
                        )
        model, experiment_args = init(config)

        env, memory, ounoise, config, summaries, saver, start_episode = experiment_args

        tic = time.time()
        monitor = run(model, experiment_args, train=True)
        monitor_test = run(model, experiment_args, train=False)

        toc = time.time()

        env.close()
        for summary in summaries:
            summary.close()
            
        monitor_macddpg_p2.append(monitor)
        monitor_macddpg_p2_test.append(monitor_test)
        
        np.save('./monitor_macddpg_p2.npy', monitor_macddpg_p2)
        np.save('./monitor_macddpg_p2_test.npy', monitor_macddpg_p2_test)
        
        print(toc-tic)<file_sep>import torch as K
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

from olbr.agents.basic import BackwardDyn

import pdb


class PPO_BD(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, 
                 clip_param, ppo_epoch, num_mini_batch, value_loss_coef, entropy_coef,
                 eps=None, max_grad_norm=None, use_clipped_value_loss=True,
                 out_func=K.sigmoid, discrete=True, object_Qfunc=None, backward_dyn=None, 
                 object_policy=None, reward_fun=None, masked_with_r=False, rnd_models=None, pred_th=0.0002,
                 dtype=K.float32, device="cuda",
                 ):

        super(PPO_BD, self).__init__()

        self.observation_space = observation_space
        self.action_space = action_space
        
        optimizer, lr = optimizer
        actor_lr, _ = lr

        self.clip_param = clip_param
        self.ppo_epoch = ppo_epoch
        self.num_mini_batch = num_mini_batch

        self.value_loss_coef = value_loss_coef
        self.entropy_coef = entropy_coef

        self.max_grad_norm = max_grad_norm
        self.use_clipped_value_loss = use_clipped_value_loss

        self.discrete = discrete

        self.object_Qfunc = object_Qfunc
        self.object_policy = object_policy

        self.dtype = dtype
        self.device = device
        self.masked_with_r = masked_with_r
        self.pred_th = pred_th

        # model initialization
        self.entities = []
        
        # actors and critics
        self.actors = []
        self.critics = []
        self.optims = []

        self.actors.append(Actor(observation_space, action_space[0], discrete, out_func).to(device))
        self.critics.append(Critic(observation_space, None).to(device))

        self.optims.append(optimizer(self.critics[0].parameters(), lr=actor_lr, eps=eps))
        self.optims[0].add_param_group({"params": self.actors[0].parameters()})
        
        self.entities.extend(self.actors)
        self.entities.extend(self.critics)
        self.entities.extend(self.optims) 

        # backward dynamics model
        if backward_dyn is not None:
            self.backward = backward_dyn
            self.backward.eval()
            self.entities.append(self.backward)

        # Learnt Q function for object
        if self.object_Qfunc is not None:
            self.object_Qfunc.eval()
            self.entities.append(self.object_Qfunc)

        # Learnt policy for object
        if self.object_policy is not None:
            self.object_policy.eval()
            self.entities.append(self.object_policy)

        if reward_fun is not None:
            self.get_obj_reward = reward_fun
        else:
            self.get_obj_reward = self.reward_fun

        if rnd_models is not None:
            self.rnd_model = rnd_models[0]
            self.rnd_target = rnd_models[1]
            self.entities.append(self.rnd_model)
            self.entities.append(self.rnd_target)
        
        print('rnd_based_coverage')

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.optims[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.optims[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False):
        with K.no_grad():
            value = self.critics[0](state.to(self.device))
            action_dist = self.actors[0](state.to(self.device))

        if exploration:
            action = action_dist.sample()
        else:
            action = action_dist.mode()

        #action = action.clamp(int(self.action_space[0].low[0]), int(self.action_space[0].high[0]))

        return value, action, action_dist.log_probs(action)

    def evaluate_actions(self, state, action, get_preactivations=False):       
        value = self.critics[0](state.to(self.device))
        action_dist = self.actors[0](state.to(self.device))

        action_log_probs = action_dist.log_probs(action)
        dist_entropy = action_dist.entropies().mean()

        if get_preactivations:
            action_preactivations = self.actors[0].get_preactivations(state.to(self.device))
        else:
            action_preactivations = None

        return value, action_log_probs, dist_entropy, action_preactivations

    def reward_fun(self, state, next_state):
        with K.no_grad():
            action = self.backward(state.to(self.device), next_state.to(self.device))
            opt_action = self.object_policy(state.to(self.device))

            reward = self.object_Qfunc(state.to(self.device), action) - self.object_Qfunc(state.to(self.device), opt_action)
        return reward.clamp(min=-1.0, max=0.0)

    def get_pred_error(self, next_state):
        with K.no_grad():
            target = self.rnd_target(next_state)
            pred = self.rnd_model(next_state)
            
        error = F.mse_loss(pred, target, reduction='none').mean(1)

        return error

    def update(self, rollouts):
        advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-5)

        value_loss_epoch = 0
        action_loss_epoch = 0
        dist_entropy_epoch = 0

        for e in range(self.ppo_epoch):
            data_generator = rollouts.feed_forward_generator(advantages, self.num_mini_batch, mini_batch_size=None)

            for sample in data_generator:
                obs_batch, actions_batch, value_preds_batch, return_batch, old_action_log_probs_batch, adv_targ = sample

                # Reshape to do in a single forward pass for all steps
                values, action_log_probs, dist_entropy, action_preactivations  = self.evaluate_actions(obs_batch, actions_batch, False)

                ratio = K.exp(action_log_probs - old_action_log_probs_batch)
                surr1 = ratio * adv_targ
                surr2 = K.clamp(ratio, 1.0 - self.clip_param, 1.0 + self.clip_param) * adv_targ
                action_loss = -K.min(surr1, surr2).mean()
                #action_loss += (action_preactivations**2).mean()*0.001

                if self.use_clipped_value_loss:
                    value_pred_clipped = value_preds_batch + (values - value_preds_batch).clamp(-self.clip_param, self.clip_param)
                    value_losses = (values - return_batch).pow(2)
                    value_losses_clipped = (value_pred_clipped - return_batch).pow(2)
                    value_loss = 0.5 * K.max(value_losses, value_losses_clipped).mean()
                else:
                    value_loss = 0.5 * (return_batch - values).pow(2).mean()

                self.optims[0].zero_grad()
                (value_loss * self.value_loss_coef + action_loss - dist_entropy * self.entropy_coef).backward()
                nn.utils.clip_grad_norm_(self.critics[0].parameters(), self.max_grad_norm)
                nn.utils.clip_grad_norm_(self.actors[0].parameters(), self.max_grad_norm)
                self.optims[0].step()

                value_loss_epoch += value_loss.item()
                action_loss_epoch += action_loss.item()
                dist_entropy_epoch += dist_entropy.item()

        num_updates = self.ppo_epoch * self.num_mini_batch

        value_loss_epoch /= num_updates
        action_loss_epoch /= num_updates
        dist_entropy_epoch /= num_updates

        return value_loss_epoch, action_loss_epoch, dist_entropy_epoch<file_sep>import numpy as np
import scipy as sc
import time
import imageio

import torch as K
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import gym_wmgds as gym

from olbr.algorithms import DDPG, DDPGC
from olbr.experience import ReplayMemory, Transition, Normalizer
from olbr.exploration import Noise, OUNoise
from olbr.utils import Saver, Summarizer, get_noise_scale, get_params, running_mean
from olbr.agents.basic import Actor 
from olbr.agents.basic import CriticReg as Critic

import pdb

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

def init(config):
    
    if config['resume'] != '':
        resume_path = config['resume']
        saver = Saver(config)
        config, start_episode, save_dict = saver.resume_ckpt()
        config['resume'] = resume_path
    else:
        start_episode = 0
        
    #hyperparameters
    ENV_NAME = config['env_id'] #'simple_spread'
    SEED = config['random_seed'] # 1

    GAMMA = config['gamma'] # 0.95
    TAU = config['tau'] # 0.01

    ACTOR_LR = config['plcy_lr'] # 0.01
    CRITIC_LR = config['crtc_lr'] # 0.01

    MEM_SIZE = config['buffer_length'] # 1000000 

    REGULARIZATION = config['regularization'] # True
    NORMALIZED_REWARDS = config['reward_normalization'] # True

    env = gym.make(ENV_NAME)
    env.seed(SEED)
    observation_space = env.observation_space.shape[0]
    action_space = env.action_space
    OUT_FUNC = K.sigmoid

    K.manual_seed(SEED)
    np.random.seed(SEED)
    
    MODEL = DDPG

    if config['verbose'] > 1:
        # utils
        summaries = (Summarizer(config['dir_summary_train'], config['port'], config['resume']),
                    Summarizer(config['dir_summary_test'], config['port'], config['resume']))
        saver = Saver(config)
    else:
        summaries = None
        saver = None

    #exploration initialization
    noise = Noise(action_space.shape[0], sigma=0.2, eps=0.3, max_u=env.action_space.high[0])
    #noise = OUNoise(action_space.shape[0])

    #model initialization
    optimizer = (optim.Adam, (ACTOR_LR, CRITIC_LR)) # optimiser func, (actor_lr, critic_lr)
    loss_func = F.mse_loss
    model = MODEL(observation_space, action_space, optimizer, 
                  Actor, Critic, loss_func, GAMMA, TAU, out_func=OUT_FUNC,
                  discrete=False, regularization=REGULARIZATION, normalized_rewards=NORMALIZED_REWARDS)
    
    if config['resume'] != '':
        for i, param in enumerate(save_dict['model_params']):
            model.entities[i].load_state_dict(param)
    
    #memory initilization
    memory = ReplayMemory(MEM_SIZE)

    experiment_args = (env, memory, noise, config, summaries, saver, start_episode)
          
    return model, experiment_args

def rollout(env, model, noise, normalizer=None, render=False):
    trajectory = []

    # monitoring variables
    episode_reward = 0
    frames = []
    
    obs = env.reset()
    nb_step = 100#env._max_episode_steps

    for i_step in range(nb_step):

        model.to_cpu()

        obs = K.tensor(obs, dtype=K.float32).unsqueeze(0)
        # Observation normalization
        if normalizer is not None:
            obs = normalizer.preprocess_with_update(obs)

        action = model.select_action(obs, noise)

        next_obs, reward, done, info = env.step(action.squeeze(0).numpy())
        reward = K.tensor(reward, dtype=dtype).view(1,1)

        # for monitoring
        episode_reward += reward
        done = 1 if (i_step == (nb_step - 1)) else 0

        trajectory.append((obs, action, reward, K.tensor(next_obs, dtype=K.float32).unsqueeze(0), done))

        # Move to the next state
        obs = next_obs

        # Record frames
        if render:
            frames.append(env.render(mode='rgb_array')[0])

    success = 0 if (info.get('is_success') is None) else info.get('is_success')

    return trajectory, episode_reward, success, frames

def run(model, experiment_args, train=True):

    total_time_start =  time.time()

    env, memory, noise, config, summaries, saver, start_episode = experiment_args

    #normalizer = Normalizer()
    
    start_episode = start_episode if train else 0
    NUM_EPISODES = config['n_episodes'] if train else config['n_episodes_test'] 
    EPISODE_LENGTH = 100#env._max_episode_steps #config['episode_length'] if train else config['episode_length_test'] 
   
    episode_reward_all = []
    episode_success_all = []
        
    for i_episode in range(start_episode, NUM_EPISODES):
        
        episode_time_start = time.time()
        #noise.scale = get_noise_scale(i_episode, config)
        if train:
            for i_cycle in range(25):
            
                trajectories = []
                for i_rollout in range(16):
                    # Initialize the environment and state
                    trajectory, _, _, _ = rollout(env, model, noise, render=0)
                    trajectories.append(trajectory)

                for trajectory in trajectories:

                    for i_step in range(len(trajectory)):
                        obs, action, reward, next_obs, done = trajectory[i_step]                    
                        if done:
                            next_obs = None

                        memory.push(obs, action, next_obs, reward)
        
            
                for _ in range(40):
                    if len(memory) > config['batch_size']-1:

                        model.to_cuda()        
                        batch = Transition(*zip(*memory.sample(config['batch_size'])))
                        critic_loss, actor_loss = model.update_parameters(batch)

                #print(normalizer.get_stats())

            # <-- end loop: i_cycle

        episode_reward_cycle = []
        episode_succeess_cycle = []
        for i_rollout in range(10):
            # Initialize the environment and state
            render = config['render'] > 0 and i_episode % config['render'] == 0
            _, episode_reward, success, frames = rollout(env, model, noise=False, normalizer=None, render=render)
                
            # Save gif
            dir_monitor = config['dir_monitor_train'] if train else config['dir_monitor_test']
            if config['render'] > 0 and i_episode % config['render'] == 0:
                imageio.mimsave('{}/{}.gif'.format(dir_monitor, i_episode), frames)

            episode_reward_cycle.append(episode_reward)
            episode_succeess_cycle.append(success)
        # <-- end loop: i_rollout 
            
        ### MONITORIRNG ###
        episode_reward_all.append(np.mean(episode_reward_cycle))
        episode_success_all.append(np.mean(episode_succeess_cycle))
                        
        if config['verbose'] > 0:
        # Printing out
            if (i_episode+1)%1 == 0:
                print("==> Episode {} of {}".format(i_episode + 1, NUM_EPISODES))
                print('  | Id exp: {}'.format(config['exp_id']))
                print('  | Exp description: {}'.format(config['exp_descr']))
                print('  | Env: {}'.format(config['env_id']))
                print('  | Process pid: {}'.format(config['process_pid']))
                print('  | Tensorboard port: {}'.format(config['port']))
                print('  | Episode total reward: {}'.format(episode_reward))
                print('  | Running mean of total reward: {}'.format(episode_reward_all[-1]))
                print('  | Success rate: {}'.format(episode_success_all[-1]))
                #print('  | Running mean of total reward: {}'.format(running_mean(episode_reward_all)[-1]))
                print('  | Time episode: {}'.format(time.time()-episode_time_start))
                print('  | Time total: {}'.format(time.time()-total_time_start))
                        
        if config['verbose'] > 0:    
            ep_save = i_episode+1 if (i_episode == NUM_EPISODES-1) else None  
            is_best_save = None
            is_best_avg_save = None   
                
            if (not train) or ((np.asarray([ep_save, is_best_save, is_best_avg_save]) == None).sum() == 3):
                to_save = False
            else:
                model.to_cpu()
                saver.save_checkpoint(save_dict   = {'model_params': [entity.state_dict() for entity in model.entities]},
                                      episode     = ep_save,
                                      is_best     = is_best_save,
                                      is_best_avg = is_best_avg_save
                                      )
                to_save = True
    
            #if (i_episode+1)%100 == 0:
            #    summary = summaries[0] if train else summaries[1]
            #    summary.update_log(i_episode, 
            #                    episode_reward, 
            #                    list(episode_reward.reshape(-1,)), 
            #                    critic_loss        = critic_loss, 
            #                    actor_loss         = actor_loss,
            #                    to_save            = to_save
            #                    )
        


            
    # <-- end loop: i_episode
    if train:
        print('Training completed')
    else:
        print('Test completed')
    
    return episode_reward_all


if __name__ == '__main__':

    monitor_macddpg_p2 = []
    monitor_macddpg_p2_test = []
    for i in range(0,5):
        config = get_params(args=['--exp_id','MACDDPG_P2_120K_'+ str(i+1), 
                                '--random_seed', str(i+1), 
                                '--agent_alg', 'MACDDPG',
                                '--protocol_type', str(2),
                                '--n_episodes', '120000',
                                '--verbose', '2',
                                ]
                        )
        model, experiment_args = init(config)

        env, memory, ounoise, config, summaries, saver, start_episode = experiment_args

        tic = time.time()
        monitor = run(model, experiment_args, train=True)
        monitor_test = run(model, experiment_args, train=False)

        toc = time.time()

        env.close()
        for summary in summaries:
            summary.close()
            
        monitor_macddpg_p2.append(monitor)
        monitor_macddpg_p2_test.append(monitor_test)
        
        np.save('./monitor_macddpg_p2.npy', monitor_macddpg_p2)
        np.save('./monitor_macddpg_p2_test.npy', monitor_macddpg_p2_test)
        
        print(toc-tic)<file_sep>import numpy as np
import scipy as sc
import time

import torch as K

from olbr.utils import get_params as get_params, running_mean, get_exp_params
from olbr.main_cher import init, run
from olbr.main_q_v2 import init as init_q
from olbr.main_q_v2 import run as run_q
from olbr.main_q_rnd import init as init_q_rnd
from olbr.main_q_rnd import run as run_q_rnd
import matplotlib.pyplot as plt

import os
import pickle
import sys
import olbr


K.set_num_threads(1)

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

exp_config = get_exp_params(sys.argv[1:])
filepath = exp_config['filepath'] 
os.chdir(filepath)

suffix = ''

if exp_config['env'] == 'PushSide':
     env_name_list = ['FetchPushObstacleSideGapMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PushMiddle':
     env_name_list = ['FetchPushObstacleMiddleGapMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PushDouble':
     env_name_list = ['FetchPushObstacleDoubleGapMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPShelf':
     env_name_list = ['FetchPickAndPlaceShelfMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPObstacle':
     env_name_list = ['FetchPickAndPlaceObstacleMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPHardest':
     env_name_list = ['FetchPickAndPlaceHardestMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPInsertion':
     env_name_list = ['FetchPickAndPlaceInsertionMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPNormal':
     env_name_list = ['FetchPickAndPlaceMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'All':
     env_name_list = ['FetchPushObstacleSideGapMulti{}-v1'.format(suffix), 
                      'FetchPushObstacleMiddleGapMulti{}-v1'.format(suffix), 
                      'FetchPushObstacleDoubleGapMulti{}-v1'.format(suffix), 
                      'FetchPickAndPlaceShelfMulti{}-v1'.format(suffix), 
                      'FetchPickAndPlaceObstacleMulti{}-v1'.format(suffix),
                      'FetchPickAndPlaceHardestMulti{}-v1'.format(suffix),
                      'FetchPickAndPlaceInsertionMulti{}-v1'.format(suffix),
                      'FetchPickAndPlaceMulti{}-v1'.format(suffix)
                      ]

for env_name in env_name_list:

    if exp_config['use_cher'] == 'True':
        use_cher = True
        print("training with CHER")
    else:
        use_cher = False
        print("training without CHER")

    if 'FetchPushObstacleDoubleGapMulti' in env_name:
        n_episodes = 200
        gamma = 0.9875
    else: 
        n_episodes = 100
        gamma = 0.98

    for i_exp in range(int(exp_config['start_n_exp']), int(exp_config['n_exp'])):
        obj_rew = False
        object_Qfunc = None
        object_policy = None  
        backward_dyn = None
        print("training without object based rewards")
        init_2 = init
        run_2 = run

        ####################### training robot ###########################  
        model_name = 'DDPG_BD'
        exp_args2=['--env_id', env_name,
                '--exp_id', model_name + '_foorob_' + str(i_exp),
                '--random_seed', str(i_exp), 
                '--agent_alg', model_name,
                '--verbose', '2',
                '--render', '0',
                '--gamma', str(gamma),
                '--n_episodes', str(n_episodes),
                '--n_cycles', '50',
                '--n_rollouts', '38',
                '--n_test_rollouts', '380',
                '--n_envs', '38',
                '--n_batches', '40',
                '--batch_size', '4864',
                '--obj_action_type', '0123456',
                '--max_nb_objects', '1',
                '--observe_obj_grp', 'False',
                ]

        config2 = get_params(args=exp_args2)
        model2, experiment_args2 = init_2(config2, agent='robot', her=use_cher, 
                                        object_Qfunc=object_Qfunc,
                                        object_policy=object_policy,
                                        backward_dyn=backward_dyn,
                                    )            
        env2, memory2, noise2, config2, normalizer2, running_rintr_mean2 = experiment_args2
        experiment_args2 = (env2, memory2, noise2, config2, normalizer2, running_rintr_mean2)

        monitor2, bestmodel = run_2(model2, experiment_args2, train=True)

        rob_name = env_name
        if use_cher:
            rob_name = rob_name + '_CHER_'
        else:
            rob_name = rob_name + '_'

        path = './ObT_models/batch1/' + rob_name + str(i_exp)
        try:  
            os.makedirs(path)
        except OSError:  
            print ("Creation of the directory %s failed" % path)
        else:  
            print ("Successfully created the directory %s" % path)

        K.save(model2.critics[0].state_dict(), path + '/robot_Qfunc.pt')
        K.save(model2.actors[0].state_dict(), path + '/robot_policy.pt')

        K.save(bestmodel[0], path + '/robot_Qfunc_best.pt')
        K.save(bestmodel[1], path + '/robot_policy_best.pt')
        
        with open(path + '/normalizer.pkl', 'wb') as file:
            pickle.dump(normalizer2, file)

        with open(path + '/normalizer_best.pkl', 'wb') as file:
            pickle.dump(bestmodel[2], file)

        path = './ObT_models/batch1/monitor_' + rob_name + str(i_exp) + '.npy'
        np.save(path, monitor2)

<file_sep>import torch as K
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

from olbr.agents.basic import BackwardDyn, RandomNetDist

import pdb


def soft_update(target, source, tau):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)

def hard_update(target, source):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(param.data)


class DDPG_BD(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, 
                 reward_fun=None, clip_Q_neg=None, nb_critics=2,
                 dtype=K.float32, device="cuda"):

        super(DDPG_BD, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        if isinstance(observation_space, (list, tuple)):
            observation_space = observation_space[0]
        self.observation_space = observation_space
        self.action_space = action_space
        self.clip_Q_neg = clip_Q_neg if clip_Q_neg is not None else -1./(1.-self.gamma)
        self.nb_critics = nb_critics

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []
        
        self.actors.append(Actor(observation_space, action_space[0], discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space, action_space[0], discrete, out_func).to(device))
        self.actors_optim.append(optimizer(self.actors[0].parameters(), lr = actor_lr))

        hard_update(self.actors_target[0], self.actors[0])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []
        
        for i_critic in range(self.nb_critics):
            self.critics.append(Critic(observation_space, action_space[0]).to(device))
            self.critics_target.append(Critic(observation_space, action_space[0]).to(device))
            self.critics_optim.append(optimizer(self.critics[i_critic].parameters(), lr = critic_lr))

        for i_critic in range(self.nb_critics):
            hard_update(self.critics_target[i_critic], self.critics[i_critic])

        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

        print('seperaate Qs for multiQ')

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False):
        self.actors[0].eval()
        with K.no_grad():
            mu = self.actors[0](state.to(self.device))
        self.actors[0].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[0].low[0]), int(self.action_space[0].high[0]))

        return mu

    def update_parameters(self, batch, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)
        
        s1 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a1 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, 0:action_space]

        s1_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        if normalizer[0] is not None:
            s1 = normalizer[0].preprocess(s1)
            s1_ = normalizer[0].preprocess(s1_)


        s, s_, a = s1, s1_, a1
        a_ = self.actors_target[0](s_)

        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device)

        for i_critic in range(self.nb_critics):
            Q = self.critics[i_critic](s, a)       
            V = self.critics_target[i_critic](s_, a_).detach()

            target_Q = (V * self.gamma) + r[:,i_critic:i_critic+1]
            target_Q = target_Q.clamp(self.clip_Q_neg, 0.)

            loss_critic = self.loss_func(Q, target_Q)

            self.critics_optim[i_critic].zero_grad()
            loss_critic.backward()
            self.critics_optim[i_critic].step()

        # actor update
        a = self.actors[0](s)

        loss_actor = 0.
        for i_critic in range(self.nb_critics):
            loss_actor += - self.critics[i_critic](s, a).mean()
        
        if self.regularization:
            loss_actor += (self.actors[0](s)**2).mean()*1

        self.actors_optim[0].zero_grad()        
        loss_actor.backward()
        self.actors_optim[0].step()
                
        return loss_critic.item(), loss_actor.item()

    def update_target(self):

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        for i_critic in range(self.nb_critics):
            soft_update(self.critics_target[i_critic], self.critics[i_critic], self.tau)



<file_sep>import torch as K
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

from olbr.agents.basic import BackwardDyn, RandomNetDist

import gym_wmgds as gym

import pdb


def soft_update(target, source, tau):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)

def hard_update(target, source):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(param.data)


class DDPG_BD(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, 
                 reward_fun=None, clip_Q_neg=None, nb_critics=2,
                 dtype=K.float32, device="cuda"):

        super(DDPG_BD, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        #if isinstance(observation_space, (list, tuple)):
        #    observation_space = observation_space[0]
        self.observation_space = observation_space
        self.action_space = action_space
        self.clip_Q_neg = clip_Q_neg if clip_Q_neg is not None else -1./(1.-self.gamma)
        self.nb_critics = nb_critics
        self.actor_action_space = gym.spaces.Box(action_space[0].low[0], 
                                                 action_space[0].high[0], 
                                                 shape=(action_space[0].shape[0]//2,), dtype='float32')

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []
        
        self.actors.append(Actor(observation_space[1], self.actor_action_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space[1], self.actor_action_space, discrete, out_func).to(device))
        self.actors_optim.append(optimizer(self.actors[0].parameters(), lr = actor_lr))

        self.actors.append(Actor(observation_space[1], self.actor_action_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space[1], self.actor_action_space, discrete, out_func).to(device))
        self.actors_optim.append(optimizer(self.actors[1].parameters(), lr = actor_lr))

        hard_update(self.actors_target[0], self.actors[0])
        hard_update(self.actors_target[1], self.actors[1])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []
        
        for i_critic in range(self.nb_critics):
            self.critics.append(Critic(observation_space[0], action_space[0]).to(device))
            self.critics_target.append(Critic(observation_space[0], action_space[0]).to(device))
            self.critics_optim.append(optimizer(self.critics[i_critic].parameters(), lr = critic_lr))

        for i_critic in range(self.nb_critics):
            hard_update(self.critics_target[i_critic], self.critics[i_critic])

        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

        print('seperaate Qs for multiQ')

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False, goal_size=None):
        if goal_size is None:
            self.actors[0].eval()
            with K.no_grad():
                mu = self.actors[0](state.to(self.device))
            self.actors[0].train()
        else:
            self.actors[0].eval()
            self.actors[1].eval()
            goal = state[:,-goal_size::]
            obs = state[:,0:-goal_size]
            obs1 = K.cat([obs[:,0:obs.shape[1]//2], goal], dim=-1)
            obs2 = K.cat([obs[:,obs.shape[1]//2::], goal], dim=-1)
            with K.no_grad():
                mu1 = self.actors[0](obs1.to(self.device))
                mu2 = self.actors[1](obs2.to(self.device))
            mu = K.cat([mu1,mu2], dim=-1)
            self.actors[0].train()
            self.actors[1].train()            
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[0].low[0]), int(self.action_space[0].high[0]))

        return mu

    def update_parameters(self, batch, normalizer=None):

        observation_space = self.observation_space[0] - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)
        
        s1 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a1 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, 0:action_space]

        s1_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        if normalizer[0] is not None:
            s1 = normalizer[0].preprocess(s1)
            s1_ = normalizer[0].preprocess(s1_)


        s, s_, a = s1, s1_, a1
        o1 = K.cat([s[:, 0:(observation_space//2)], s[:, observation_space::]], dim=-1)
        o2 = K.cat([s[:, (observation_space//2):observation_space], s[:, observation_space::]], dim=-1)

        o = [o1, o2]

        o1_ = K.cat([s_[:, 0:observation_space//2], s_[:, observation_space::]], dim=-1)
        o2_ = K.cat([s_[:, observation_space//2:observation_space], s_[:, observation_space::]], dim=-1)

        a1_ = self.actors_target[0](o1_)
        a2_ = self.actors_target[1](o2_)

        a_ = K.cat([a1_, a2_],dim=-1)

        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device)

        for i_critic in range(self.nb_critics):
            Q = self.critics[i_critic](s, a)       
            V = self.critics_target[i_critic](s_, a_).detach()

            target_Q = (V * self.gamma) + r[:,i_critic:i_critic+1]
            target_Q = target_Q.clamp(self.clip_Q_neg, 0.)

            loss_critic = self.loss_func(Q, target_Q)

            self.critics_optim[i_critic].zero_grad()
            loss_critic.backward()
            self.critics_optim[i_critic].step()

        # actor update
        for i_actor in range(2):
            a1 = self.actors[0](o[0])
            a2 = self.actors[1](o[1])

            a = K.cat([a1, a2],dim=-1)

            loss_actor = 0.
            for i_critic in range(self.nb_critics):
                loss_actor += - self.critics[i_critic](s, a).mean()
            
            if self.regularization:
                loss_actor += (self.actors[i_actor](o[i_actor])**2).mean()*1

            self.actors_optim[i_actor].zero_grad()        
            loss_actor.backward()
            self.actors_optim[i_actor].step()
                
        return loss_critic.item(), loss_actor.item()

    def update_target(self):

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.actors_target[1], self.actors[1], self.tau)
        for i_critic in range(self.nb_critics):
            soft_update(self.critics_target[i_critic], self.critics[i_critic], self.tau)



<file_sep>import numpy as np
import scipy as sc
import time
import pickle
import os
import sys

import torch as K
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from olbr.utils import get_params as get_params, running_mean, get_exp_params
from olbr.main_rnd import init, run, interfere
from olbr.agents.basic import TrajectoryDyn

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

K.set_num_threads(1)

class MyDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data[0])
    
    def __getitem__(self,idx):
        return self.data[0][idx], self.data[1][idx], self.data[2][idx], self.data[3][idx]

exp_config = get_exp_params(sys.argv[1:])
filepath = exp_config['filepath'] 
os.chdir(filepath)

suffix = ''
if exp_config['env'] == 'PushSide':
     env_name_list = ['FetchPushObstacleSideGapMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PushMiddle':
     env_name_list = ['FetchPushObstacleMiddleGapMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PushDouble':
     env_name_list = ['FetchPushObstacleDoubleGapMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPShelf':
     env_name_list = ['FetchPickAndPlaceShelfMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPObstacle':
     env_name_list = ['FetchPickAndPlaceObstacleMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPHardest':
     env_name_list = ['FetchPickAndPlaceHardestMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPInsertion':
     env_name_list = ['FetchPickAndPlaceInsertionMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'PnPNormal':
     env_name_list = ['FetchPickAndPlaceMulti{}-v1'.format(suffix)]
elif exp_config['env'] == 'All':
     env_name_list = ['FetchPushObstacleSideGapMulti{}-v1'.format(suffix), 
                      'FetchPushObstacleMiddleGapMulti{}-v1'.format(suffix), 
                      'FetchPushObstacleDoubleGapMulti{}-v1'.format(suffix), 
                      'FetchPickAndPlaceShelfMulti{}-v1'.format(suffix), 
                      'FetchPickAndPlaceObstacleMulti{}-v1'.format(suffix),
                      'FetchPickAndPlaceHardestMulti{}-v1'.format(suffix),
                      'FetchPickAndPlaceInsertionMulti{}-v1'.format(suffix),
                      'FetchPickAndPlaceMulti{}-v1'.format(suffix)
                      ]

for env_name in env_name_list:

    if 'FetchPushObstacleDoubleGapMulti' in env_name:
        n_episodes = 50
        gamma = 0.9875
    else: 
        n_episodes = 25
        gamma = 0.98

    if env_name == 'FetchPushObstacleSideGapMulti-v1':
        path = './ObT_models/obj/push_side_7d_ep25/'
    elif env_name == 'FetchPushObstacleMiddleGapMulti-v1':
        path = './ObT_models/obj/push_middle_7d_ep25/'
    elif env_name == 'FetchPushObstacleDoubleGapMulti-v1':
        path = './ObT_models/obj/push_double_7d_ep50/'
    elif env_name == 'FetchPickAndPlaceShelfMulti-v1':
        path = './ObT_models/obj/pnp_shelf_7d_ep25/'
    elif env_name == 'FetchPickAndPlaceObstacleMulti-v1':
        path = './ObT_models/obj/pnp_obstacle_7d_ep25/'
    elif env_name == 'FetchPickAndPlaceHardestMulti-v1':
        path = './ObT_models/obj/pnp_hardest_7d_ep25/'
    elif env_name == 'FetchPickAndPlaceInsertionMulti-v1':
        path = './ObT_models/obj/pnp_insertion_7d_ep25/'
    print('object trajectory model will be created in')
    print(path)        


    #training object policy
    model_name = 'DDPG_BD'
    exp_args=['--env_id', env_name,
            '--exp_id', model_name + '_foorob_' + str(0),
            '--random_seed', str(0), 
            '--agent_alg', model_name,
            '--verbose', '2',
            '--render', '0',
            '--gamma', str(gamma),
            '--n_episodes', str(n_episodes),
            '--n_cycles', '50',
            '--n_rollouts', '38',
            '--n_test_rollouts', '380',
            '--n_envs', '38',
            '--n_batches', '40',
            '--batch_size', '4864',
            '--obj_action_type', '0123456',
            '--max_nb_objects', '1',
            '--observe_obj_grp', 'False',
            '--rob_policy', '02',
            '--buffer_length', '10000000',
            ]

    config = get_params(args=exp_args)
    model, experiment_args = init(config, agent='object', her=True, 
                                object_Qfunc=None, 
                                backward_dyn=None,
                                object_policy=None)
    env, memory, noise, config, normalizer, agent_id = experiment_args
    monitor = run(model, experiment_args, train=True)

    try:  
        os.makedirs(path)
    except OSError:  
        print ("Creation of the directory %s failed" % path)
    else:  
        print ("Successfully created the directory %s" % path)

    K.save(model.critics[0].state_dict(), path + 'object_Qfunc.pt')
    K.save(model.actors[0].state_dict(), path + 'object_policy.pt')
    K.save(model.backward.state_dict(), path + 'backward_dyn.pt')
    with open(path + 'normalizer.pkl', 'wb') as file:
        pickle.dump(normalizer, file)  
        
    np.save(path + 'monitor', monitor)

    #generating object trajectories
    memory.clear_buffer()
    interfere(model, experiment_args)

    #creating the dataset from trajectories
    time_step = 1
    horizon=n_episodes
    goal_len = memory.buffers['ag'].shape[2]

    ag_in_all = []
    ag_out_all = []
    g_in_all = []
    step_in_all = []
    for i_time in range(time_step,horizon+time_step,time_step):
        ag_in = memory.buffers['ag'][0:memory.current_size,0,0:goal_len]
        if i_time == horizon:
            ag_out = memory.buffers['g'][0:memory.current_size,0,0:goal_len]
        else:
            ag_out = memory.buffers['ag'][0:memory.current_size,i_time,0:goal_len]
        g_in = memory.buffers['g'][0:memory.current_size,0,0:goal_len]
        
        ag_in_all.append(ag_in.reshape(-1,goal_len).astype('float32'))
        ag_out_all.append(ag_out.reshape(-1,goal_len).astype('float32'))
        g_in_all.append(g_in.reshape(-1,goal_len).astype('float32'))
        step_in_all.append((np.ones_like(ag_in)[:,0]*i_time).astype('float32'))
        
    ag_in_all = np.asarray(ag_in_all).reshape(-1,goal_len)
    ag_out_all = np.asarray(ag_out_all).reshape(-1,goal_len)
    g_in_all = np.asarray(g_in_all).reshape(-1,goal_len)
    step_in_all = np.asarray(step_in_all).reshape(-1,1)
        
    ag_in_mean = memory.buffers['ag'][0:memory.current_size].mean((0,1))
    ag_in_std = memory.buffers['ag'][0:memory.current_size].std((0,1))
    g_in_mean = memory.buffers['g'][0:memory.current_size].mean((0,1))
    g_in_std = memory.buffers['g'][0:memory.current_size].std((0,1))

    step_in_all_mean = step_in_all.mean(0)
    step_in_all_std = step_in_all.std(0)

    obj_mean = [ag_in_mean, g_in_mean, step_in_all_mean]
    obj_std = [ag_in_std, g_in_std, step_in_all_std]

    with open(path + 'objGoal_mean.pkl', 'wb') as file:
        pickle.dump(obj_mean, file)
    with open(path + 'objGoal_std.pkl', 'wb') as file:
        pickle.dump(obj_std, file)
        
    train_split = 0.7
    train_bound = int(train_split*len(ag_in_all))
    idxs = np.random.permutation(len(ag_in_all))
    train_idxs = idxs[0:train_bound]
    test_idxs = idxs[train_bound::]

    ag_in_all -= ag_in_mean
    ag_in_all /= (ag_in_std + 1e-5)

    g_in_all -= g_in_mean
    g_in_all /= (g_in_std + 1e-5)

    step_in_all -= step_in_all_mean
    step_in_all /= step_in_all_std

    train_data = (ag_in_all[train_idxs], g_in_all[train_idxs], step_in_all[train_idxs], ag_out_all[train_idxs])
    test_data = (ag_in_all[test_idxs], g_in_all[test_idxs], step_in_all[test_idxs], ag_out_all[test_idxs])

    train_dataset = MyDataset(train_data)
    test_dataset = MyDataset(test_data)

    train_loader = DataLoader(train_dataset, batch_size=128)
    test_loader = DataLoader(test_dataset, batch_size=128)

    #initialising the trajectory model
    model_objtraj = TrajectoryDyn(goal_len).to(device)
    optimizer = optim.Adam(model_objtraj.parameters())

    #training the trajectory model
    nb_traj_epochs = 30
    for epoch in range(nb_traj_epochs): 

        # training
        running_loss = 0.0
        total = 0.

        model_objtraj.train()
        for data in train_loader:
            # get the inputs
            X1, X2, X3, T = data
            X1 = X1.to(device)
            X2 = X2.to(device)
            X3 = X3.to(device)
            T = T.to(device)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            Y = model_objtraj(X1, X2, X3)
            loss = F.mse_loss(Y, T)
            loss.backward()
            optimizer.step()

            # print statistics
            total += T.size(0)
            running_loss += loss.item()#*T.size(0)

        #running_loss /= total

        # testing
        correct_Y = 0.
        total_test = 0. 

        model_objtraj.eval()
        with K.no_grad():
            for data in test_loader:
                # get the inputs
                X1, X2, X3, T = data
                X1 = X1.to(device)
                X2 = X2.to(device)
                X3 = X3.to(device)
                T = T.to(device)
                # 
                Y = model_objtraj(X1, X2, X3)

                total_test += T.size(0)
                loss = F.mse_loss(Y, T)
                
                correct_Y += loss.item()
                
            #correct_Y /= total_test

        print('%d, loss: %.3f, acc = %.3f %%'  % (epoch + 1, 
                                                running_loss, 
                                                correct_Y))
    print('Finished Training of Object Trajectory Model')                                            
    K.save(model_objtraj.state_dict(), path + 'objGoal_model.pt')<file_sep>import numpy as np
import scipy as sc
import time

import torch as K

from olbr.utils import get_params as get_params, running_mean, get_exp_params
from olbr.main import init, run
from olbr.main_ppo import init as init_ppo
from olbr.main_ppo import run as run_ppo
import matplotlib.pyplot as plt

import os
import pickle
import sys

filepath='/jmain01/home/JAD022/grm01/oxk28-grm01/Dropbox/Jupyter/notebooks/Reinforcement_Learning/'
os.chdir(filepath)

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

exp_config = get_exp_params(sys.argv[1:])

if exp_config['env'] == 'Push':
    env_name = 'FetchPushMulti-v1'
elif exp_config['env'] == 'PnP':
    env_name = 'FetchPickAndPlaceMulti-v1'

if exp_config['multiseed'] == 'True':
    multiseed = True
    n_envs = 38
    from olbr.main import init, run
elif exp_config['multiseed'] == 'False':
    multiseed = False
    n_envs = 1
    from olbr.main import init, run

masked_with_r = exp_config['masked_with_r']
if exp_config['use_her'] == 'True':
    use_her = True
else:
    use_her = False

for i_exp in range(0,int(exp_config['n_exp'])):
    if exp_config['obj_rew'] == 'True':
    ####################### loading object ###########################

        model_name = 'DDPG_BD'
        exp_args=['--env_id', env_name,
                '--exp_id', model_name + '_fooobj_' + str(0),
                '--random_seed', str(0), 
                '--agent_alg', model_name,
                '--verbose', '2',
                '--render', '0',
                '--gamma', '0.98',
                '--n_episodes', '20',
                '--n_cycles', '50',
                '--n_rollouts', '38',
                '--n_test_rollouts', '38',
                '--n_envs', str(n_envs),
                '--n_batches', '40',
                '--batch_size', '256',
                '--n_bd_batches', '400',
                '--obj_action_type', '012',
                '--max_nb_objects', '1',
                '--observe_obj_grp', 'False',
                '--rob_policy', '01',
                ]

        config = get_params(args=exp_args)
        model, experiment_args = init(config, agent='object', her=True, 
                                    object_Qfunc=None, 
                                    backward_dyn=None,
                                    object_policy=None)
        env, memory, noise, config, normalizer, agent_id = experiment_args

        #loading the object model
        if exp_config['env'] == 'Push':
            path = './models/obj/obj_push_xyz/'
        elif exp_config['env'] == 'PnP':
            path = './models/obj/obj_pnp_xyz/'

        model.critics[0].load_state_dict(K.load(path + 'object_Qfunc.pt'))
        model.backward.load_state_dict(K.load(path + 'backward_dyn.pt'))
        model.actors[0].load_state_dict(K.load(path + 'object_policy.pt'))
        with open(path + 'normalizer.pkl', 'rb') as file:
            normalizer = pickle.load(file)

        experiment_args = (env, memory, noise, config, normalizer, agent_id)
        
        obj_rew = True
        object_Qfunc = model.critics[0]
        backward_dyn = model.backward
        object_policy = model.actors[0]  
    ####################### loading object ###########################
    elif exp_config['obj_rew'] == 'False':
        obj_rew = False
        object_Qfunc = None
        backward_dyn = None
        object_policy = None  

    ####################### training robot ###########################  
    if exp_config['rob_model'] == 'DDPG':
        model_name = 'DDPG_BD'
        exp_args2=['--env_id', env_name,
                '--exp_id', model_name + '_foorob_' + str(i_exp),
                '--random_seed', str(i_exp), 
                '--agent_alg', model_name,
                '--verbose', '2',
                '--render', '0',
                '--gamma', '0.98',
                '--n_episodes', '50',
                '--n_cycles', '50',
                '--n_rollouts', '38',
                '--n_test_rollouts', '380',
                '--n_envs', str(n_envs),
                '--n_batches', '40',
                '--batch_size', '256',
                '--n_bd_batches', '400',
                '--obj_action_type', '012',
                '--max_nb_objects', '1',
                '--observe_obj_grp', 'False',
                '--rob_policy', '01',
                '--masked_with_r', masked_with_r,
                ]

        config2 = get_params(args=exp_args2)
        model2, experiment_args2 = init(config2, agent='robot', her=use_her, 
                                        object_Qfunc=object_Qfunc, 
                                        backward_dyn=backward_dyn,
                                        object_policy=object_policy
                                    )
        env2, memory2, noise2, config2, normalizer2, agent_id2 = experiment_args2
        if obj_rew:
            normalizer2[1] = normalizer[1]
        experiment_args2 = (env2, memory2, noise2, config2, normalizer2, agent_id2)

        monitor2 = run(model2, experiment_args2, train=True)

    elif exp_config['rob_model'] == 'PPO':
        model_name = 'PPO_BD'
        exp_args2=['--env_id', env_name,
                '--exp_id', model_name + '_foorob_' + str(i_exp),
                '--random_seed', str(i_exp), 
                '--agent_alg', model_name,
                '--verbose', '2',
                '--render', '0',
                '--gamma', '0.98',
                '--n_episodes', '50',
                '--n_cycles', '50',
                '--n_rollouts', '38',
                '--n_batches', '4',
                '--batch_size', '256',
                '--reward_normalization', 'False', 
                '--ai_object_rate', '0.0',
                '--obj_action_type', '012',
                '--max_nb_objects', '1',
                '--observe_obj_grp', 'True',
                '--masked_with_r', masked_with_r,
                '--n_test_rollouts', '100',
                '--plcy_lr', '3e-4',
                '--crtc_lr', '3e-4',
                '--ppo_epoch', '3',
                '--entropy_coef', '0.00',
                '--clip_param', '0.1',
                '--use_gae', "False"
        ]

        config2 = get_params(args=exp_args2)
        model2, experiment_args2 = init_ppo(config2, agent='robot', her=False, 
                                        object_Qfunc=object_Qfunc, 
                                        backward_dyn=backward_dyn,
                                        object_policy=object_policy
                                    )
        env2, memory2, noise2, config2, normalizer2, agent_id2 = experiment_args2
        normalizer2[1] = normalizer[1]
        experiment_args2 = (env2, memory2, noise2, config2, normalizer2, agent_id2)

        monitor2 = run_ppo(model2, experiment_args2, train=True)

    rob_name = exp_config['env']
    if exp_config['rob_model'] == 'PPO':
        if obj_rew:
            rob_name = rob_name + '_v4_PPO_'
        else:
            rob_name = rob_name + '_PPO_'
    elif exp_config['rob_model'] == 'DDPG':
        if obj_rew:
            if use_her:
                rob_name = rob_name + '_v4_HER_'
            else:
                rob_name = rob_name + '_v4_DDPG_'
        else:
            if use_her:
                rob_name = rob_name + '_HER_'
            else:
                rob_name = rob_name + '_DDPG_'

    rob_name = rob_name + 'Norm_Slide_Clipped_Both_'
    if masked_with_r:
        rob_name = rob_name + 'Masked_PlusR'
    else:
        rob_name = rob_name + 'PlusR'

    if multiseed:
        rob_name = rob_name + '_Multiseed'
    else:
        rob_name = rob_name + '_Singleseed'

    path = './models/_recent/rob_model_' + rob_name + '_' + str(i_exp)
    try:  
        os.makedirs(path)
    except OSError:  
        print ("Creation of the directory %s failed" % path)
    else:  
        print ("Successfully created the directory %s" % path)

    K.save(model2.critics[0].state_dict(), path + '/robot_Qfunc.pt')
    K.save(model2.actors[0].state_dict(), path + '/robot_policy.pt')
    if obj_rew:
        K.save(model2.object_Qfunc.state_dict(), path + '/object_Qfunc.pt')
        K.save(model2.backward.state_dict(), path + '/backward_dyn.pt')
        K.save(model2.object_policy.state_dict(), path + '/object_policy.pt')
    
    with open(path + '/normalizer.pkl', 'wb') as file:
        pickle.dump(normalizer2, file)

    path = './monitors/_recent/monitor_' + rob_name  + '_' + str(i_exp) + '.npy'
    np.save(path, monitor2)

<file_sep># Object Locomotion based Rewards

<file_sep>import torch as K
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

from olbr.exploration import gumbel_softmax
from olbr.agents.basic import ForwardDynReg, AutoEncReg, AutoEncNextReg

import pdb


def soft_update(target, source, tau):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)

def hard_update(target, source):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(param.data)


class DDPG(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(DDPG, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []
        
        self.actors.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.actors_optim.append(optimizer(self.actors[0].parameters(), lr = actor_lr))

        hard_update(self.actors_target[0], self.actors[0])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []
        
        self.critics.append(Critic(observation_space, action_space).to(device))
        self.critics_target.append(Critic(observation_space, action_space).to(device))
        self.critics_optim.append(optimizer(self.critics[0].parameters(), lr = critic_lr))

        hard_update(self.critics_target[0], self.critics[0])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False):
        self.actors[0].eval()
        with K.no_grad():
            mu = self.actors[0](state.to(self.device))
        self.actors[0].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space.low[0]), int(self.action_space.high[0]))

        return mu

    def update_parameters(self, batch, normalizer=None):
        
        mask = K.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch.state), 1), device=self.device)

        s = K.cat(batch.state, dim=0).to(self.device)
        a = K.cat(batch.action, dim=0).to(self.device)
        r = K.cat(batch.reward, dim=0).to(self.device)
        s_ = K.cat([i.to(self.device) for i in batch.next_state if i is not None], dim=0)
        a_ = K.zeros_like(a)[:,0:s_.shape[0],]
        
        if normalizer[0] is not None:
            s = normalizer[0].preprocess(s)
            s_ = normalizer[0].preprocess(s_)

        if self.normalized_rewards:
            if normalizer[1] is not None:
                r = r
                #r = normalizer[1].preprocess(r)
            else:
                r -= r.mean()
                r /= r.std()
        
        Q = self.critics[0](s, a)        
        a_ = self.actors_target[0](s_)
        V[mask] = self.critics_target[0](s_, a_).detach()

        #loss_critic = self.loss_func(Q, (V * self.gamma) + r)
        target_Q = (V * self.gamma) + r
        target_Q = target_Q.clamp(-1./(1.-self.gamma), 0.)
        
        loss_critic = self.loss_func(Q, target_Q)

        self.critics_optim[0].zero_grad()
        loss_critic.backward()
        K.nn.utils.clip_grad_norm_(self.critics[0].parameters(), 0.5)
        self.critics_optim[0].step()

        a = self.actors[0](s)

        loss_actor = -self.critics[0](s, a).mean()
        
        if self.regularization:
            #loss_actor += (self.actors[0].get_preactivations(s)**2).mean()*1
            loss_actor += (self.actors[0](s)**2).mean()*1

        self.actors_optim[0].zero_grad()        
        loss_actor.backward()
        K.nn.utils.clip_grad_norm_(self.actors[0].parameters(), 0.5)
        self.actors_optim[0].step()

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)

        return loss_critic.item(), loss_actor.item()


class DDPG_GAR(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(DDPG_GAR, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []
        
        self.actors.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.actors_optim.append(optimizer(self.actors[0].parameters(), lr = actor_lr))

        hard_update(self.actors_target[0], self.actors[0])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []
        
        self.critics.append(Critic(observation_space, action_space).to(device))
        self.critics_target.append(Critic(observation_space, action_space).to(device))
        self.critics_optim.append(optimizer(self.critics[0].parameters(), lr = critic_lr))

        hard_update(self.critics_target[0], self.critics[0])

        self.critics.append(Critic(observation_space, action_space).to(device))
        self.critics_target.append(Critic(observation_space, action_space).to(device))
        self.critics_optim.append(optimizer(self.critics[1].parameters(), lr = critic_lr))

        hard_update(self.critics_target[1], self.critics[1])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

        # Forward Dynamics Network
        self.forward = AutoEncNextReg(observation_space, action_space).to(device)
        #self.forward = AutoEncReg(observation_space, action_space).to(device)
        #self.forward = ForwardDynReg(observation_space, action_space).to(device)
        self.forward_optim = optimizer(self.forward.parameters(), lr = critic_lr)

        self.entities.append(self.forward)
        self.entities.append(self.forward_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False):
        self.actors[0].eval()
        with K.no_grad():
            mu = self.actors[0](state.to(self.device))
        self.actors[0].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space.low[0]), int(self.action_space.high[0]))

        return mu

    def update_parameters(self, batch, normalizer=None):
        
        mask = K.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch.state), 1), device=self.device)

        s = K.cat(batch.state, dim=0).to(self.device)
        a = K.cat(batch.action, dim=0).to(self.device)
        r = K.cat(batch.reward, dim=0).to(self.device)
        s_ = K.cat([i.to(self.device) for i in batch.next_state if i is not None], dim=0)
        a_ = K.zeros_like(a)[:,0:s_.shape[0],]

        if self.normalized_rewards:
            if normalizer[1] is not None:
                r = r
                #r = normalizer[1].preprocess(r)
            else:
                r -= r.mean()
                r /= r.std()
        
        if normalizer[0] is not None:
            normed_s = normalizer[0].preprocess(s)
            normed_s_ = normalizer[0].preprocess(s_)

        #p = self.forward.get_intr_rewards(s,a)
        p = K.zeros_like(r)
        p[mask], M = self.forward.get_intr_rewards(normed_s_)
        if self.normalized_rewards:
            if normalizer[2] is not None:
                p = normalizer[2].preprocess_with_update(p.to('cpu')).to(self.device)          
            else:
                p -= p.mean()
                p /= p.std()
        
        # extrinsic Q update
        Q = self.critics[0](normed_s, a)

        a_ = self.actors_target[0](normed_s_)

        V[mask] = self.critics_target[0](normed_s_, a_).detach()

        loss_critic = self.loss_func(Q, (V * self.gamma) + r + 0.1*p)

        self.critics_optim[0].zero_grad()
        loss_critic.backward()
        K.nn.utils.clip_grad_norm_(self.critics[0].parameters(), 0.5)
        self.critics_optim[0].step()

        # intrinsic Q update
        Q = self.critics[1](normed_s, a)

        V[mask] = self.critics_target[1](normed_s_, a_).detach()

        loss_critic = self.loss_func(Q, (V * self.gamma) + p)
        # in case of reward propagation
        #reg_critic =  self.critics[0].GAR(1, 1, 0.000001)
        #loss_critic += reg_critic

        self.critics_optim[1].zero_grad()
        loss_critic.backward()
        K.nn.utils.clip_grad_norm_(self.critics[1].parameters(), 0.5)
        self.critics_optim[1].step()

        # actor update
        a = self.actors[0](normed_s)

        loss_actor = - self.critics[0](normed_s, a).mean() - 0.1*self.critics[1](normed_s, a).mean() #- self.critics[0].GAR(1, 1, 0.000001)
        
        if self.regularization:
            loss_actor += (self.actors[0].get_preactivations(normed_s)**2).mean()*1

        self.actors_optim[0].zero_grad()        
        loss_actor.backward()
        K.nn.utils.clip_grad_norm_(self.actors[0].parameters(), 0.5)
        self.actors_optim[0].step()

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)
        soft_update(self.critics_target[1], self.critics[1], self.tau)

        p = p.squeeze(1).cpu().numpy()
        r = r.squeeze(1).cpu().numpy()

        ind = np.argsort(p)

        return loss_critic.item(), loss_actor.item(), p[ind], r[ind], M

    def update_forward(self, batch, normalizer=None):
        mask = K.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=K.uint8, device=self.device)

        s = K.cat(batch.state, dim=0).to(self.device)
        a = K.cat(batch.action, dim=0).to(self.device)
        s_ = K.cat([i.to(self.device) for i in batch.next_state if i is not None], dim=0)
        
        if normalizer[0] is not None:
            s = normalizer[0].preprocess(s)
            s_ = normalizer[0].preprocess(s_)

        #pred = self.forward(s,a)
        pred = self.forward(s_)
        #loss_forward = self.loss_func(pred[mask], s_)
        loss_forward = self.loss_func(pred, s_)
        #loss_forward = self.loss_func(pred, K.cat([s, a], dim=1))
        reg_forward = self.forward.GAR(0.03, 0.03, 0.001, False)
        #reg_forward = self.forward.GAR(1, 1, 1)
        loss_forward += reg_forward

        self.forward_optim.zero_grad()
        loss_forward.backward()
        #K.nn.utils.clip_grad_norm_(self.forward.parameters(), 0.5)
        self.forward_optim.step()

        return loss_forward.item()-reg_forward.item(), reg_forward.item()


class HDDPG(object):
    def __init__(self, observation_space, action_space, goal_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(HDDPG, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.action_space = action_space
        self.goal_space = goal_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []
        
        self.actors.append(Actor(observation_space+goal_space.shape[0], action_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space+goal_space.shape[0], action_space, discrete, out_func).to(device))

        self.actors.append(Actor(observation_space, goal_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space, goal_space, discrete, out_func).to(device))

        for i in range(len(self.actors)):
            self.actors_optim.append(optimizer(self.actors[i].parameters(), lr = actor_lr))
            hard_update(self.actors_target[i], self.actors[i])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []
        
        self.critics.append(Critic(observation_space+goal_space.shape[0], action_space).to(device))
        self.critics_target.append(Critic(observation_space+goal_space.shape[0], action_space).to(device))

        self.critics.append(Critic(observation_space, goal_space).to(device))
        self.critics_target.append(Critic(observation_space, goal_space).to(device))

        for i in range(len(self.critics)):
            self.critics_optim.append(optimizer(self.critics[i].parameters(), lr = critic_lr))
            hard_update(self.critics_target[i], self.critics[i])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'   

    def select_goal(self, state, exploration=False):
        self.actors[1].eval()
        with K.no_grad():
            mu = self.actors[1](state.to(self.device))
        self.actors[1].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)

        #mu = mu * K.tensor(self.goal_space.high[0], dtype=self.dtype, device=self.device)
        mu = mu.clamp(-1, 1)

        return mu

    def select_action(self, state, exploration=False):
        self.actors[0].eval()
        with K.no_grad():
            mu = self.actors[0](state.to(self.device))
        self.actors[0].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        
        mu = mu * K.tensor(self.action_space.high[0], dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space.low[0]), int(self.action_space.high[0]))

        return mu

    def update_parameters(self, batch):

        for i in range(2):
            
            mask = K.tensor(tuple(map(lambda s: s is not None, batch[i].next_state)), dtype=K.uint8, device=self.device)

            V = K.zeros((len(batch[i].state), 1), device=self.device)

            s = K.cat(batch[i].state, dim=0).to(self.device)
            a = K.cat(batch[i].action, dim=0).to(self.device)
            r = K.cat(batch[i].reward, dim=0).to(self.device)
            s_ = K.cat([i.to(self.device) for i in batch[i].next_state if i is not None], dim=0)
            a_ = K.zeros_like(a)[:,0:s_.shape[0],]

            if self.normalized_rewards:
                r -= r.mean()
                r /= r.std()
            
            Q = self.critics[i](s, a)
            
            a_ = self.actors_target[i](s_)

            V[mask] = self.critics_target[i](s_, a_).detach()

            loss_critic = self.loss_func(Q, (V * self.gamma) + r) 

            self.critics_optim[i].zero_grad()
            loss_critic.backward()
            K.nn.utils.clip_grad_norm_(self.critics[i].parameters(), 0.5)
            self.critics_optim[i].step()

            a = self.actors[i](s)

            loss_actor = -self.critics[i](s, a).mean()
            
            if self.regularization:
                loss_actor += (self.actors[i].get_preactivations(s)**2).mean()*1

            self.actors_optim[i].zero_grad()        
            loss_actor.backward()
            K.nn.utils.clip_grad_norm_(self.actors[i].parameters(), 0.5)
            self.actors_optim[i].step()

            soft_update(self.actors_target[i], self.actors[i], self.tau)
            soft_update(self.critics_target[i], self.critics[i], self.tau)

        
        return loss_critic.item(), loss_actor.item()


class DDPGC(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(DDPGC, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []
        
        self.actors.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.actors_target.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.actors_optim.append(optimizer(self.actors[0].parameters(), lr = actor_lr))

        hard_update(self.actors_target[0], self.actors[0])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []
        
        self.critics.append(Critic(observation_space, action_space).to(device))
        self.critics_target.append(Critic(observation_space, action_space).to(device))
        self.critics_optim.append(optimizer(self.critics[0].parameters(), lr = critic_lr))

        hard_update(self.critics_target[0], self.critics[0])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)


        # actors
        self.cactors = []
        self.cactors_target = []
        self.cactors_optim = []
        
        self.cactors.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.cactors_target.append(Actor(observation_space, action_space, discrete, out_func).to(device))
        self.cactors_optim.append(optimizer(self.cactors[0].parameters(), lr = actor_lr))

        hard_update(self.cactors_target[0], self.cactors[0])

        self.entities.extend(self.cactors)
        self.entities.extend(self.cactors_target)
        self.entities.extend(self.cactors_optim) 
        
        # critics   
        self.ccritics = []
        self.ccritics_target = []
        self.ccritics_optim = []
        
        self.ccritics.append(Critic(observation_space, action_space).to(device))
        self.ccritics_target.append(Critic(observation_space, action_space).to(device))
        self.ccritics_optim.append(optimizer(self.ccritics[0].parameters(), lr = critic_lr))

        hard_update(self.ccritics_target[0], self.ccritics[0])
            
        self.entities.extend(self.ccritics)
        self.entities.extend(self.ccritics_target)
        self.entities.extend(self.ccritics_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, exploration=False):
        self.actors[0].eval()
        self.cactors[0].eval()
        with K.no_grad():
            mu = self.actors[0](state.to(self.device))
            cmu = self.cactors[0](state.to(self.device))
        self.actors[0].train()
        self.cactors[0].train()

        mu = mu*(1 - cmu)

        mu = K.tensor(self.action_space.low[0], dtype=mu.dtype, device=mu.device) + mu*K.tensor((self.action_space.high[0]-self.action_space.low[0]), dtype=mu.dtype, device=mu.device)
        
        if exploration:
            mu += K.tensor(exploration.noise(), dtype=self.dtype, device=self.device)
            cmu += K.tensor(exploration.noise(), dtype=self.dtype, device=self.device)

        mu = mu.clamp(int(self.action_space.low[0]), int(self.action_space.high[0]))

        return mu

    def update_parameters(self, batch):
        
        mask = K.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch.state), 1), device=self.device)

        s = K.cat(batch.state, dim=0).to(self.device)
        a = K.cat(batch.action, dim=0).to(self.device)
        r = K.cat(batch.reward, dim=0).to(self.device)
        s_ = K.cat([i.to(self.device) for i in batch.next_state if i is not None], dim=0)
        a_ = K.zeros_like(a)[:,0:s_.shape[0],]

        if self.normalized_rewards:
            r -= r.mean()
            r /= r.std()
        
        Q = self.critics[0](s, a)
        
        a_ = self.actors_target[0](s_)

        V[mask] = self.critics_target[0](s_, a_).detach()

        loss_critic = self.loss_func(Q, (V * self.gamma) + r) 

        self.critics_optim[0].zero_grad()
        loss_critic.backward()
        K.nn.utils.clip_grad_norm_(self.critics[0].parameters(), 0.5)
        self.critics_optim[0].step()

        a = self.actors[0](s)

        loss_actor = -self.critics[0](s, a).mean()
        
        if self.regularization:
            loss_actor += (self.actors[0].get_preactivations(s)**2).mean()*1e-3

        self.actors_optim[0].zero_grad()        
        loss_actor.backward()
        K.nn.utils.clip_grad_norm_(self.actors[0].parameters(), 0.5)
        self.actors_optim[0].step()

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)


        mask = K.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch.state), 1), device=self.device)

        s = K.cat(batch.state, dim=0).to(self.device)
        a = K.cat(batch.action, dim=0).to(self.device)
        r = K.cat(batch.reward, dim=0).to(self.device)
        s_ = K.cat([i.to(self.device) for i in batch.next_state if i is not None], dim=0)
        a_ = K.zeros_like(a)[:,0:s_.shape[0],]

        if self.normalized_rewards:
            r -= r.mean()
            r /= r.std()

        Q = self.ccritics[0](s, a)
        
        a_ = self.cactors_target[0](s_)

        V[mask] = self.ccritics_target[0](s_, a_).detach()

        loss_critic = self.loss_func(Q, (V * self.gamma) + (1-r)) 

        self.ccritics_optim[0].zero_grad()
        loss_critic.backward()
        K.nn.utils.clip_grad_norm_(self.ccritics[0].parameters(), 0.5)
        self.ccritics_optim[0].step()

        a = self.cactors[0](s)

        loss_actor = -self.ccritics[0](s, a).mean()
        
        if self.regularization:
            loss_actor += (self.cactors[0].get_preactivations(s)**2).mean()*1e-3

        self.cactors_optim[0].zero_grad()        
        loss_actor.backward()
        K.nn.utils.clip_grad_norm_(self.cactors[0].parameters(), 0.5)
        self.cactors_optim[0].step()

        soft_update(self.cactors_target[0], self.cactors[0], self.tau)
        soft_update(self.ccritics_target[0], self.ccritics[0], self.tau)

        
        return loss_critic.item(), loss_actor.item()


# class MAHCDDPG(PARENT):
#     def __init__(self, num_agents, observation_space, action_space, medium_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
#                  discrete=True, regularization=False, normalized_rewards=False, communication=None, discrete_comm=True, Comm_Actor=None, Comm_Critic=None, dtype=K.float32, device="cuda"):
        
#         super().__init__(num_agents, observation_space, action_space, medium_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func,
#                          discrete, regularization, normalized_rewards, communication, discrete_comm, Comm_Actor, Comm_Critic, dtype, device)

#         optimizer, lr = optimizer
#         actor_lr, critic_lr, comm_actor_lr, comm_critic_lr = lr

#         # model initialization
#         self.entities = []
        
#         # actors
#         self.actors = []
#         self.actors_target = []
#         self.actors_optim = []
        
#         for i in range(num_agents):
#             self.actors.append(Actor(observation_space+medium_space, action_space, discrete, out_func).to(device))
#             self.actors_target.append(Actor(observation_space+medium_space, action_space, discrete, out_func).to(device))
#             self.actors_optim.append(optimizer(self.actors[i].parameters(), actor_lr))

#         for i in range(num_agents):
#             hard_update(self.actors_target[i], self.actors[i])

#         self.entities.extend(self.actors)
#         self.entities.extend(self.actors_target)
#         self.entities.extend(self.actors_optim) 
        
#         # critics   
#         self.critics = []
#         self.critics_target = []
#         self.critics_optim = []
        
#         for i in range(num_agents):
#             self.critics.append(Critic(observation_space+medium_space, action_space).to(device))
#             self.critics_target.append(Critic(observation_space+medium_space, action_space).to(device))
#             self.critics_optim.append(optimizer(self.critics[i].parameters(), lr = critic_lr))

#         for i in range(num_agents):
#             hard_update(self.critics_target[i], self.critics[i])
            
#         self.entities.extend(self.critics)
#         self.entities.extend(self.critics_target)
#         self.entities.extend(self.critics_optim)    

#         # communication actors
#         self.comm_actors = []
#         self.comm_actors_target = []
#         self.comm_actors_optim = []
        
#         for i in range(num_agents):
#             self.comm_actors.append(Comm_Actor(observation_space, 1, discrete_comm, K.sigmoid).to(device))
#             self.comm_actors_target.append(Comm_Actor(observation_space, 1, discrete_comm, K.sigmoid).to(device))
#             self.comm_actors_optim.append(optimizer(self.comm_actors[i].parameters(), lr = comm_actor_lr))
            
#         for i in range(num_agents):
#             hard_update(self.comm_actors_target[i], self.comm_actors[i])

#         self.entities.extend(self.comm_actors)
#         self.entities.extend(self.comm_actors_target)
#         self.entities.extend(self.comm_actors_optim)

#         # communication critics   
#         self.comm_critics = []
#         self.comm_critics_target = []
#         self.comm_critics_optim = []
        
#         for i in range(num_agents):
#             self.comm_critics.append(Comm_Critic(observation_space*num_agents, 1*num_agents).to(device))
#             self.comm_critics_target.append(Comm_Critic(observation_space*num_agents, 1*num_agents).to(device))
#             self.comm_critics_optim.append(optimizer(self.comm_critics[i].parameters(), lr = comm_critic_lr))

#         for i in range(num_agents):
#             hard_update(self.comm_critics_target[i], self.comm_critics[i]) 

#         print('amanin')

#     def select_comm_action(self, state, i_agent, exploration=False):
#         self.comm_actors[i_agent].eval()
#         with K.no_grad():
#             mu = self.comm_actors[i_agent](state.to(self.device))
#         self.actors[i_agent].train()
#         if self.discrete_comm:
#             mu = gumbel_softmax(mu, exploration=exploration)
#         else:
#             if exploration:
#                 mu += K.tensor(exploration.noise(), dtype=self.dtype, device=self.device)
#                 mu = mu.clamp(0, 1) 
#         return mu

#     def update_parameters(self, batch, batch2, i_agent):
        
#         mask = K.tensor(tuple(map(lambda s: s is not None, batch.next_state)), dtype=K.uint8, device=self.device)

#         V = K.zeros((len(batch.state), 1), device=self.device)

#         s = K.cat(batch.state, dim=1).to(self.device)
#         a = K.cat(batch.action, dim=1).to(self.device)
#         r = K.cat(batch.reward, dim=1).to(self.device)
#         s_ = K.cat([i.to(self.device) for i in batch.next_state if i is not None], dim=1)
#         a_ = K.zeros_like(a)[:,0:s_.shape[1],]
        
#         m = K.cat(batch.medium, dim=1).to(self.device)

#         if self.normalized_rewards:
#             r -= r.mean()
#             r /= r.std()

#         Q = self.critics[i_agent](K.cat([s[[i_agent],], m], dim=-1),
#                                   a[[i_agent],])
        
#         for i in range(self.num_agents):
#             a_[i,] = self.actors_target[i](K.cat([s_[[i],], m[:,mask,]], dim=-1))

#         V[mask] = self.critics_target[i_agent](K.cat([s_[[i_agent],], m[:,mask,]], dim=-1),
#                                                a_[[i_agent],]).detach()

#         loss_critic = self.loss_func(Q, (V * self.gamma) + r[[i_agent],].squeeze(0)) 

#         self.critics_optim[i_agent].zero_grad()
#         loss_critic.backward()
#         K.nn.utils.clip_grad_norm_(self.critics[i_agent].parameters(), 0.5)
#         self.critics_optim[i_agent].step()

#         for i in range(self.num_agents):
#             a[i,] = self.actors[i](K.cat([s[[i],], m], dim=-1))

#         loss_actor = -self.critics[i_agent](K.cat([s[[i_agent],], m], dim=-1), 
#                                             a[[i_agent],]).mean()
        
#         if self.regularization:
#             loss_actor += (self.actors[i_agent].get_preactivations(K.cat([s[[i_agent],], m], dim=-1))**2).mean()*1e-3

#         self.actors_optim[i_agent].zero_grad()        
#         loss_actor.backward()
#         K.nn.utils.clip_grad_norm_(self.actors[i_agent].parameters(), 0.5)
#         self.actors_optim[i_agent].step()

#         soft_update(self.actors_target[i_agent], self.actors[i_agent], self.tau)
#         soft_update(self.critics_target[i_agent], self.critics[i_agent], self.tau)

#         ## update of the communication part

#         mask = K.tensor(tuple(map(lambda s: s is not None, batch2.next_state)), dtype=K.uint8, device=self.device)
#         V = K.zeros((len(batch2.state), 1), device=self.device)

#         s = K.cat(batch2.state, dim=1).to(self.device)
#         r = K.cat(batch2.reward, dim=1).to(self.device)
#         s_ = K.cat([i.to(self.device) for i in batch2.next_state if i is not None], dim=1)
        
#         c = K.cat(batch2.comm_action, dim=1).to(self.device)
#         c_ = K.zeros_like(c)[:,0:s_.shape[1],]
        
#         if self.normalized_rewards:
#             r -= r.mean()
#             r /= r.std()

#         Q = self.comm_critics[i_agent](s, c)

#         for i in range(self.num_agents):
#             if self.discrete_comm:
#                 c_[i,] = gumbel_softmax(self.comm_actors_target[i](s_[[i],]), exploration=False)
#             else:
#                 c_[i,] = self.comm_actors_target[i](s_[[i],])

#         V[mask] = self.comm_critics_target[i_agent](s_, c_).detach()

#         loss_critic = self.loss_func(Q, (V * self.gamma) + r[[i_agent],].squeeze(0)) 

#         self.comm_critics_optim[i_agent].zero_grad()
#         loss_critic.backward()
#         K.nn.utils.clip_grad_norm_(self.comm_critics[i_agent].parameters(), 0.5)
#         self.comm_critics_optim[i_agent].step()

#         for i in range(self.num_agents):
#             if self.discrete_comm:
#                 c[i,] = gumbel_softmax(self.comm_actors[i](s[[i],]), exploration=False)
#             else:
#                 c[i,] = self.comm_actors[i](s[[i],])

#         loss_actor = -self.comm_critics[i_agent](s, c).mean()
        
#         if self.regularization:
#             loss_actor += (self.comm_actors[i_agent].get_preactivations(s[[i_agent],])**2).mean()*1e-3

#         self.comm_actors_optim[i_agent].zero_grad()        
#         loss_actor.backward()
#         K.nn.utils.clip_grad_norm_(self.comm_actors[i_agent].parameters(), 0.5)
#         self.comm_actors_optim[i_agent].step()

#         soft_update(self.comm_actors_target[i_agent], self.comm_actors[i_agent], self.tau)
#         soft_update(self.comm_critics_target[i_agent], self.comm_critics[i_agent], self.tau)

        
#         return loss_critic.item(), loss_actor.item()
<file_sep>import torch as K
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np

from olbr.agents.basic import BackwardDyn

import pdb


def soft_update(target, source, tau):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)

def hard_update(target, source):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(param.data)


class MADDPG(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(MADDPG, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.observation_space = observation_space
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []

        for i in range(2):
            self.actors.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_target.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_optim.append(optimizer(self.actors[i].parameters(), lr = actor_lr))

        for i in range(2):
            hard_update(self.actors_target[i], self.actors[i])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []

        for i in range(2):
            self.critics.append(Critic(observation_space*2, action_space[2]).to(device))
            self.critics_target.append(Critic(observation_space*2, action_space[2]).to(device))
            self.critics_optim.append(optimizer(self.critics[i].parameters(), lr = critic_lr))
        
        for i in range(2):
            hard_update(self.critics_target[i], self.critics[i])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, i_agent, exploration=False):
        self.actors[i_agent].eval()
        with K.no_grad():
            mu = self.actors[i_agent](state.to(self.device))
        self.actors[i_agent].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[i_agent].low[0]), int(self.action_space[i_agent].high[0]))

        return mu

    def update_parameters(self, batch, i_agent, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        mask = K.tensor(tuple(map(lambda ai_object: ai_object==0, K.tensor(batch['o'][:,-1]))), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)

        s1 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a1 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, 0:action_space]
        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device).unsqueeze(1)

        s1_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)


        if normalizer[0] is not None:
            s1 = normalizer[0].preprocess(s1)
            s1_ = normalizer[0].preprocess(s1_)

        if normalizer[1] is not None:
            s2 = normalizer[1].preprocess(s2)
            s2_ = normalizer[1].preprocess(s2_)

        s = [s1, s2]

        a1_ = self.actors_target[0](s1_)
        a2_ = self.actors_target[1](s2_)
        a2_[mask] *= 0.00
        
        # Critics
        # updating the critic of the robot
        Q = self.critics[i_agent](K.cat([s1,s2],dim=1), K.cat([a1,a2],dim=1))     
        V = self.critics_target[i_agent](K.cat([s1_,s2_],dim=1), K.cat([a1_,a2_],dim=1)).detach()

        target_Q = (V * self.gamma) + r
        target_Q = target_Q.clamp(-1./(1.-self.gamma), 0.)
        loss_critic = self.loss_func(Q, target_Q)

        self.critics_optim[i_agent].zero_grad()
        loss_critic.backward()
        self.critics_optim[i_agent].step()

        # Actors 
        a1 = self.actors[0](s1)
        a2 = self.actors[1](s2)
        a2[mask] *= 0.00

        loss_actor = -self.critics[i_agent](K.cat([s1,s2],dim=1), K.cat([a1,a2],dim=1)).mean()
        
        if self.regularization:
            loss_actor += (self.actors[i_agent](s[i_agent])**2).mean()*1

        self.actors_optim[i_agent].zero_grad()        
        loss_actor.backward()
        #K.nn.utils.clip_grad_norm_(self.actors[0].parameters(), 0.5)
        self.actors_optim[i_agent].step()

        return loss_critic.item(), loss_actor.item()

    def update_target(self):

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)

        soft_update(self.actors_target[1], self.actors[1], self.tau)
        soft_update(self.critics_target[1], self.critics[1], self.tau)


class DDPG(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(DDPG, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.observation_space = observation_space
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []

        for i in range(2):
            self.actors.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_target.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_optim.append(optimizer(self.actors[i].parameters(), lr = actor_lr))

        for i in range(2):
            hard_update(self.actors_target[i], self.actors[i])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []

        for i in range(2):
            self.critics.append(Critic(observation_space, action_space[i]).to(device))
            self.critics_target.append(Critic(observation_space, action_space[i]).to(device))
            self.critics_optim.append(optimizer(self.critics[i].parameters(), lr = critic_lr))
        
        for i in range(2):
            hard_update(self.critics_target[i], self.critics[i])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, i_agent, exploration=False):
        self.actors[i_agent].eval()
        with K.no_grad():
            mu = self.actors[i_agent](state.to(self.device))
        self.actors[i_agent].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[i_agent].low[0]), int(self.action_space[i_agent].high[0]))

        return mu

    def update_parameters(self, batch, i_agent, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        mask = K.tensor(tuple(map(lambda ai_object: ai_object==0, K.tensor(batch['o'][:,-1]))), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)

        s1 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a1 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, 0:action_space]
        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device).unsqueeze(1)

        s1_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)


        if normalizer[0] is not None:
            s1 = normalizer[0].preprocess(s1)
            s1_ = normalizer[0].preprocess(s1_)

        if normalizer[1] is not None:
            s2 = normalizer[1].preprocess(s2)
            s2_ = normalizer[1].preprocess(s2_)

        a1_ = self.actors_target[0](s1_)
        a2_ = self.actors_target[1](s2_)
        a2_[mask] *= 0.00            

        s = [s1, s2]
        s_ = [s1_, s2_]
        a = [a1, a2]
        a_ = [a1_, a2_]
        
        # Critics
        Q = self.critics[i_agent](s[i_agent], a[i_agent])     
        V = self.critics_target[i_agent](s_[i_agent], a_[i_agent]).detach()

        target_Q = (V * self.gamma) + r
        target_Q = target_Q.clamp(-1./(1.-self.gamma), 0.)
        loss_critic = self.loss_func(Q, target_Q)

        self.critics_optim[i_agent].zero_grad()
        loss_critic.backward()
        self.critics_optim[i_agent].step()

        # Actors 
        a1 = self.actors[0](s1)
        a2 = self.actors[1](s2)
        a2[mask] *= 0.00
        a = [a1, a2]

        loss_actor = -self.critics[i_agent](s[i_agent], a[i_agent]).mean()
        
        if self.regularization:
            loss_actor += (self.actors[i_agent](s[i_agent])**2).mean()*1

        self.actors_optim[i_agent].zero_grad()        
        loss_actor.backward()
        self.actors_optim[i_agent].step()

        return loss_critic.item(), loss_actor.item()

    def update_target(self):

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)

        soft_update(self.actors_target[1], self.actors[1], self.tau)
        soft_update(self.critics_target[1], self.critics[1], self.tau)


class MADDPG_R(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(MADDPG_R, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.observation_space = observation_space
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []

        for i in range(2):
            self.actors.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_target.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_optim.append(optimizer(self.actors[i].parameters(), lr = actor_lr))

        for i in range(2):
            hard_update(self.actors_target[i], self.actors[i])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []

        for i in range(2):
            self.critics.append(Critic(observation_space, action_space[2]).to(device))
            self.critics_target.append(Critic(observation_space, action_space[2]).to(device))
            self.critics_optim.append(optimizer(self.critics[i].parameters(), lr = critic_lr))
        
        for i in range(2):
            hard_update(self.critics_target[i], self.critics[i])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, i_agent, exploration=False):
        self.actors[i_agent].eval()
        with K.no_grad():
            mu = self.actors[i_agent](state.to(self.device))
        self.actors[i_agent].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[i_agent].low[0]), int(self.action_space[i_agent].high[0]))

        return mu

    def update_parameters(self, batch, i_agent, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        mask = K.tensor(tuple(map(lambda ai_object: ai_object==0, K.tensor(batch['o'][:,-1]))), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)

        s1 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a1 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, 0:action_space]
        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device).unsqueeze(1)

        s1_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)


        if normalizer[0] is not None:
            s1 = normalizer[0].preprocess(s1)
            s1_ = normalizer[0].preprocess(s1_)

        if normalizer[1] is not None:
            s2 = normalizer[1].preprocess(s2)
            s2_ = normalizer[1].preprocess(s2_)

        a1_ = self.actors_target[0](s1_)
        a2_ = self.actors_target[1](s2_)
        a2_[mask] *= 0.00            

        s = [s1, s2]
        s_ = [s1_, s2_]
        
        # Critics
        Q = self.critics[i_agent](s[i_agent], K.cat([a1, a2],dim=1))     
        V = self.critics_target[i_agent](s_[i_agent], K.cat([a1_, a2_],dim=1)).detach()

        target_Q = (V * self.gamma) + r
        target_Q = target_Q.clamp(-1./(1.-self.gamma), 0.)
        loss_critic = self.loss_func(Q, target_Q)

        self.critics_optim[i_agent].zero_grad()
        loss_critic.backward()
        self.critics_optim[i_agent].step()

        # Actors 
        a1 = self.actors[0](s1)
        a2 = self.actors[1](s2)
        a2[mask] *= 0.00

        loss_actor = -self.critics[i_agent](s[i_agent], K.cat([a1, a2],dim=1)).mean()
        
        if self.regularization:
            loss_actor += (self.actors[i_agent](s[i_agent])**2).mean()*1

        self.actors_optim[i_agent].zero_grad()        
        loss_actor.backward()
        self.actors_optim[i_agent].step()

        return loss_critic.item(), loss_actor.item()

    def update_target(self):

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)

        soft_update(self.actors_target[1], self.actors[1], self.tau)
        soft_update(self.critics_target[1], self.critics[1], self.tau)


class MADDPG_RAE(object):
    def __init__(self, observation_space, action_space, optimizer, Actor, Critic, loss_func, gamma, tau, out_func=K.sigmoid,
                 discrete=True, regularization=False, normalized_rewards=False, dtype=K.float32, device="cuda"):

        super(MADDPG_RAE, self).__init__()

        optimizer, lr = optimizer
        actor_lr, critic_lr = lr

        self.loss_func = loss_func
        self.gamma = gamma
        self.tau = tau
        self.out_func = out_func
        self.discrete = discrete
        self.regularization = regularization
        self.normalized_rewards = normalized_rewards
        self.dtype = dtype
        self.device = device
        self.observation_space = observation_space
        self.action_space = action_space

        # model initialization
        self.entities = []
        
        # actors
        self.actors = []
        self.actors_target = []
        self.actors_optim = []

        for i in range(2):
            self.actors.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_target.append(Actor(observation_space, action_space[i], discrete, out_func).to(device))
            self.actors_optim.append(optimizer(self.actors[i].parameters(), lr = actor_lr))

        for i in range(2):
            hard_update(self.actors_target[i], self.actors[i])

        self.entities.extend(self.actors)
        self.entities.extend(self.actors_target)
        self.entities.extend(self.actors_optim) 
        
        # critics   
        self.critics = []
        self.critics_target = []
        self.critics_optim = []

        for i in range(2):
            self.critics.append(Critic(observation_space, action_space[2]).to(device))
            self.critics_target.append(Critic(observation_space, action_space[2]).to(device))
            self.critics_optim.append(optimizer(self.critics[i].parameters(), lr = critic_lr))
        
        for i in range(2):
            hard_update(self.critics_target[i], self.critics[i])
            
        self.entities.extend(self.critics)
        self.entities.extend(self.critics_target)
        self.entities.extend(self.critics_optim)

        # backward dynamics model
        self.backward = BackwardDyn(3, action_space[1]).to(device)
        self.backward_optim = optimizer(self.backward.parameters(), lr = critic_lr)

        self.entities.append(self.backward)
        self.entities.append(self.backward_optim)

    def to_cpu(self):
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cpu()
        self.device = 'cpu'

    def to_cuda(self):        
        for entity in self.entities:
            if type(entity) != type(self.actors_optim[0]):
                entity.cuda()
        self.device = 'cuda'    

    def select_action(self, state, i_agent, exploration=False):
        self.actors[i_agent].eval()
        with K.no_grad():
            mu = self.actors[i_agent](state.to(self.device))
        self.actors[i_agent].train()
        if exploration:
            mu = K.tensor(exploration.get_noisy_action(mu.cpu().numpy()), dtype=self.dtype, device=self.device)
        mu = mu.clamp(int(self.action_space[i_agent].low[0]), int(self.action_space[i_agent].high[0]))

        return mu

    def update_parameters(self, batch, i_agent, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        mask = K.tensor(tuple(map(lambda ai_object: ai_object==0, K.tensor(batch['o'][:,-1]))), dtype=K.uint8, device=self.device)

        V = K.zeros((len(batch['o']), 1), dtype=self.dtype, device=self.device)

        s1 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
                    K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a1 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, 0:action_space]
        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        r = K.tensor(batch['r'], dtype=self.dtype, device=self.device).unsqueeze(1)

        s1_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, 0:observation_space],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)
        s2__ = K.cat([K.tensor(batch['o_3'], dtype=self.dtype, device=self.device)[:, observation_space:],
                     K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        ag = K.tensor(batch['ag'], dtype=self.dtype, device=self.device)
        ag_2 = K.tensor(batch['ag_2'], dtype=self.dtype, device=self.device)  
        ag_3 = K.tensor(batch['ag_3'], dtype=self.dtype, device=self.device)             


        if normalizer[0] is not None:
            s1 = normalizer[0].preprocess(s1)
            s1_ = normalizer[0].preprocess(s1_)

        if normalizer[1] is not None:
            s2 = normalizer[1].preprocess(s2)
            s2_ = normalizer[1].preprocess(s2_)
            s2__ = normalizer[1].preprocess(s2__)

        a1_ = self.actors_target[0](s1_)
        a2_ = self.actors_target[1](s2_)
        #a2_[mask] = self.estimate_obj_action(s2_[mask], s2__[mask])   
        a2_[mask] = self.estimate_obj_action(ag_2[mask], ag_3[mask])         

        s = [s1, s2]
        s_ = [s1_, s2_]
        
        # Critics
        Q = self.critics[i_agent](s[i_agent], K.cat([a1, a2],dim=1))     
        V = self.critics_target[i_agent](s_[i_agent], K.cat([a1_, a2_],dim=1)).detach()

        target_Q = (V * self.gamma) + r
        target_Q = target_Q.clamp(-1./(1.-self.gamma), 0.)
        loss_critic = self.loss_func(Q, target_Q)

        self.critics_optim[i_agent].zero_grad()
        loss_critic.backward()
        self.critics_optim[i_agent].step()

        # Actors 
        a1 = self.actors[0](s1)
        a2 = self.actors[1](s2)
        #a2_[mask] = self.estimate_obj_action(s2[mask], s2_[mask])
        a2[mask] = self.estimate_obj_action(ag[mask], ag_2[mask])

        loss_actor = -self.critics[i_agent](s[i_agent], K.cat([a1, a2],dim=1)).mean()
        
        if self.regularization:
            loss_actor += (self.actors[i_agent](s[i_agent])**2).mean()*1

        self.actors_optim[i_agent].zero_grad()        
        loss_actor.backward()
        self.actors_optim[i_agent].step()

        return loss_critic.item(), loss_actor.item()

    def update_target(self):

        soft_update(self.actors_target[0], self.actors[0], self.tau)
        soft_update(self.critics_target[0], self.critics[0], self.tau)

        soft_update(self.actors_target[1], self.actors[1], self.tau)
        soft_update(self.critics_target[1], self.critics[1], self.tau)

    def estimate_obj_action(self, state2, next_state2):
        self.backward.eval()
        with K.no_grad():
            action2 = self.backward(state2.to(self.device), next_state2.to(self.device))
        self.backward.train()

        return action2

    def update_backward(self, batch, normalizer=None):

        observation_space = self.observation_space - K.tensor(batch['g'], dtype=self.dtype, device=self.device).shape[1]
        action_space = self.action_space[0].shape[0]

        mask = K.tensor(tuple(map(lambda ai_object: ai_object>0, K.tensor(batch['o'][:,-1]))), dtype=K.uint8, device=self.device)

        #s2 = K.cat([K.tensor(batch['o'], dtype=self.dtype, device=self.device)[:, observation_space:],
        #            K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        a2 = K.tensor(batch['u'], dtype=self.dtype, device=self.device)[:, action_space:]

        #s2_ = K.cat([K.tensor(batch['o_2'], dtype=self.dtype, device=self.device)[:, observation_space:],
        #             K.tensor(batch['g'], dtype=self.dtype, device=self.device)], dim=-1)

        #if normalizer[1] is not None:
        #    s2 = normalizer[1].preprocess(s2)
        #    s2_ = normalizer[1].preprocess(s2_)

        #a2_pred = self.backward(s2[mask], s2_[mask])

        ag = K.tensor(batch['ag'], dtype=self.dtype, device=self.device)
        ag_2 = K.tensor(batch['ag_2'], dtype=self.dtype, device=self.device)

        a2_pred = self.backward(ag[mask], ag_2[mask])

        loss_backward = self.loss_func(a2_pred, a2[mask])

        self.backward_optim.zero_grad()
        loss_backward.backward()
        #K.nn.utils.clip_grad_norm_(self.forward.parameters(), 0.5)
        self.backward_optim.step()

        return loss_backward.item()




<file_sep>import numpy as np
import scipy as sc
import time
import imageio

import torch as K
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import gym_wmgds as gym

from olbr.algorithms.ppo_fine import PPO_BD
from olbr.experience import Normalizer
from olbr.exploration import Noise
from olbr.utils import Saver, Summarizer, get_params, running_mean
from olbr.agents.basic import ActorStoch as Actor 
from olbr.agents.basic import Critic

import pdb

import matplotlib
import matplotlib.pyplot as plt

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

def init(config, agent='robot', her=False, object_Qfunc=None, backward_dyn=None, object_policy=None, reward_fun=None):
        
    #hyperparameters
    ENV_NAME = config['env_id'] 
    SEED = config['random_seed']
    N_ENVS = config['n_envs']

    env = []
    if 'Fetch' in ENV_NAME and 'Multi' in ENV_NAME:
        for i_env in range(N_ENVS):
            env.append(gym.make(ENV_NAME, n_objects=config['max_nb_objects'], 
                                          obj_action_type=config['obj_action_type'], 
                                          observe_obj_grp=config['observe_obj_grp'],
                                          obj_range=config['obj_range']))
        n_rob_actions = 4
        n_actions = config['max_nb_objects'] * len(config['obj_action_type']) + n_rob_actions
    elif 'HandManipulate' in ENV_NAME and 'Multi' in ENV_NAME:
        for i_env in range(N_ENVS):
            env.append(gym.make(ENV_NAME, obj_action_type=config['obj_action_type']))
        n_rob_actions = 20
        n_actions = 1 * len(config['obj_action_type']) + n_rob_actions
    else:
        for i_env in range(N_ENVS):
            env.append(gym.make(ENV_NAME))

    def her_reward_fun(ag_2, g, info):  # vectorized
        return env[0].compute_reward(achieved_goal=ag_2, desired_goal=g, info=info)
    
    for i_env in range(N_ENVS):
        env[i_env].seed(SEED+10*i_env)
    K.manual_seed(SEED)
    np.random.seed(SEED)

    observation_space = env[0].observation_space.spaces['observation'].shape[1] + env[0].observation_space.spaces['desired_goal'].shape[0]
    action_space = (gym.spaces.Box(-1., 1., shape=(n_rob_actions,), dtype='float32'),
                    gym.spaces.Box(-1., 1., shape=(n_actions-n_rob_actions,), dtype='float32'),
                    gym.spaces.Box(-1., 1., shape=(n_actions,), dtype='float32'))

    GAMMA = config['gamma'] 
    TAU = config['tau'] 
    ACTOR_LR = config['plcy_lr'] 
    CRITIC_LR = config['crtc_lr'] 

    MEM_SIZE = config['buffer_length']

    REGULARIZATION = config['regularization']
    NORMALIZED_REWARDS = config['reward_normalization']

    if config['agent_alg'] == 'PPO_BD':
        MODEL = PPO_BD
        OUT_FUNC = 'linear'
        from olbr.replay_buffer import ReplayBuffer
        from olbr.her_sampler import make_sample_her_transitions
        from olbr.replay_buffer import RolloutStorage as RolloutStorage

    #exploration initialization
    env[0]._max_episode_steps *= config['max_nb_objects']
    noise = (True,
             Noise(action_space[1].shape[0], sigma=0.2, eps=0.3)
            )

    #model initialization
    optimizer = (optim.Adam, (ACTOR_LR, CRITIC_LR)) # optimiser func, (actor_lr, critic_lr)
    loss_func = F.mse_loss           
    model = MODEL(observation_space, action_space, optimizer, Actor, Critic, 
            config['clip_param'], config['ppo_epoch'], config['n_batches'], config['value_loss_coef'], config['entropy_coef'],
            eps=config['eps'], max_grad_norm=config['max_grad_norm'], use_clipped_value_loss=True,
            out_func=OUT_FUNC, discrete=False, 
            agent_id=0, object_Qfunc=object_Qfunc, backward_dyn=backward_dyn, 
            object_policy=object_policy, reward_fun=reward_fun, masked_with_r=config['masked_with_r'])
    normalizer = [Normalizer(), Normalizer()]

    #memory initilization  
    if her:
        sample_her_transitions = make_sample_her_transitions('future', 4, her_reward_fun)
    else:
        sample_her_transitions = make_sample_her_transitions('none', 4, her_reward_fun)

    buffer_shapes = {
        'o' : (env[0]._max_episode_steps, env[0].observation_space.spaces['observation'].shape[1]*2),
        'ag' : (env[0]._max_episode_steps, env[0].observation_space.spaces['achieved_goal'].shape[0]),
        'g' : (env[0]._max_episode_steps, env[0].observation_space.spaces['desired_goal'].shape[0]),
        'u' : (env[0]._max_episode_steps-1, action_space[2].shape[0])
        }
    memory = (RolloutStorage(env[0]._max_episode_steps-1, config['n_rollouts'], (observation_space,), action_space[0]),
              ReplayBuffer(buffer_shapes, MEM_SIZE, env[0]._max_episode_steps, sample_her_transitions)
            )

    experiment_args = (env, memory, noise, config, normalizer, 0)
          
    return model, experiment_args

def rollout(env, model, noise, i_env, normalizer=None, render=False, agent_id=0, ai_object=False, rob_policy=[0., 0.]):
    trajectories = []
    for i_agent in range(2):
        trajectories.append([])
    
    # monitoring variables
    episode_reward = 0
    frames = []
    
    env[i_env].env.ai_object = ai_object
    env[i_env].env.deactivate_ai_object() 
    state_all = env[i_env].reset()

    if ai_object:
        for i_step in range(env[0]._max_episode_steps):
            model.to_cpu()

            obs = [K.tensor(obs, dtype=K.float32).unsqueeze(0) for obs in state_all['observation']]
            goal = K.tensor(state_all['desired_goal'], dtype=K.float32).unsqueeze(0)

            # Observation normalization
            obs_goal = []
            obs_goal.append(K.cat([obs[0], goal], dim=-1))
            if normalizer[0] is not None:
                obs_goal[0] = normalizer[0].preprocess(obs_goal[0])

            value, action, log_probs = model.select_action(obs_goal[0], noise[0])
            value = value.cpu().numpy().squeeze(0)
            action = action.cpu().numpy().squeeze(0)
            log_probs = log_probs.cpu().numpy().squeeze(0)

            action_to_env = np.zeros_like(env[0].action_space.sample())
            action_to_env[0:action.shape[0]] = action

            next_state_all, reward, done, info = env[i_env].step(action_to_env)

            # Move to the next state
            state_all = next_state_all

            # Record frames
            if render:
                frames.append(env[i_env].render(mode='rgb_array')[0])

        env[i_env].env.activate_ai_object() 
    
    for i_step in range(env[0]._max_episode_steps):

        model.to_cpu()

        obs = [K.tensor(obs, dtype=K.float32).unsqueeze(0) for obs in state_all['observation']]
        goal = K.tensor(state_all['desired_goal'], dtype=K.float32).unsqueeze(0)

        # Observation normalization
        obs_goal = []
        for i_agent in range(2):
            obs_goal.append(K.cat([obs[i_agent], goal], dim=-1))
            if normalizer[i_agent] is not None:
                obs_goal[i_agent] = normalizer[i_agent].preprocess_with_update(obs_goal[i_agent])

        value, action, log_probs = model.select_action(obs_goal[0], noise[0])
        value = value.cpu().numpy().squeeze(0)
        action = action.cpu().numpy().squeeze(0)
        log_probs = log_probs.cpu().numpy().squeeze(0)

        if ai_object:
            action_to_env = env[0].action_space.sample() * rob_policy[0] + np.ones_like(env[0].action_space.sample()) * rob_policy[1]
            action_to_env[action.shape[0]::] = model.get_obj_action(obs_goal[1], noise[1]).cpu().numpy().squeeze(0)
        else:
            action_to_env = np.zeros_like(env[0].action_space.sample())
            action_to_env[0:action.shape[0]] = action

        next_state_all, reward, done, info = env[i_env].step(action_to_env)

        if ai_object:
            for i_agent in range(2):
                state = {
                    'observation'   : state_all['observation'][i_agent],
                    'achieved_goal' : state_all['achieved_goal'],
                    'desired_goal'  : state_all['desired_goal']   
                    }
                next_state = {
                    'observation'   : next_state_all['observation'][i_agent],
                    'achieved_goal' : next_state_all['achieved_goal'],
                    'desired_goal'  : next_state_all['desired_goal']    
                    }
                
                trajectories[i_agent].append((state.copy(), action_to_env, reward, next_state.copy(), done))
        else:
            next_obs = [K.tensor(next_obs, dtype=K.float32).unsqueeze(0) for next_obs in next_state_all['observation']]
            
            # Observation normalization
            next_obs_goal = []
            for i_agent in range(2):
                next_obs_goal.append(K.cat([next_obs[i_agent], goal], dim=-1))
                if normalizer[i_agent] is not None:
                    next_obs_goal[i_agent] = normalizer[i_agent].preprocess(next_obs_goal[i_agent])

            if model.masked_with_r:
                reward = model.get_obj_reward(obs_goal[1], next_obs_goal[1]).cpu().numpy().squeeze(0) * np.abs(reward) + (reward)
            else:
                reward = model.get_obj_reward(obs_goal[1], next_obs_goal[1]).cpu().numpy().squeeze(0) + (reward)

            # for monitoring
            episode_reward += reward

            for i_agent in range(2):
                trajectories[i_agent].append((obs_goal[i_agent].cpu().numpy().squeeze(0), 
                                                action_to_env[0:action.shape[0]], 
                                                reward, 
                                                next_obs_goal[i_agent].cpu().numpy().squeeze(0), 
                                                done, 
                                                value, 
                                                log_probs))

        # Move to the next state
        state_all = next_state_all

        # Record frames
        if render:
            frames.append(env[i_env].render(mode='rgb_array')[0])

    if ai_object:
        obs, ags, goals, acts = [], [], [], []

        for trajectory in trajectories:
            obs.append([])
            ags.append([])
            goals.append([])
            acts.append([])
            for i_step in range(env[0]._max_episode_steps):
                obs[-1].append(trajectory[i_step][0]['observation'])
                ags[-1].append(trajectory[i_step][0]['achieved_goal'])
                goals[-1].append(trajectory[i_step][0]['desired_goal'])
                if (i_step < env[0]._max_episode_steps - 1): 
                    acts[-1].append(trajectory[i_step][1])

        trajectories = {
            'o'    : np.concatenate(obs,axis=1)[np.newaxis,:],
            'ag'   : np.asarray(ags)[0:1,],
            'g'    : np.asarray(goals)[0:1,],
            'u'    : np.asarray(acts)[0:1,],
        }
    else:
        obs, acts, rews, vals, logs = [], [], [], [], []

        for i_step in range(env[0]._max_episode_steps):
            obs.append(trajectories[0][i_step][0])
            vals.append(trajectories[0][i_step][5])    
            if (i_step < env[0]._max_episode_steps - 1): 
                acts.append(trajectories[0][i_step][1])
                rews.append(trajectories[0][i_step][2])
                logs.append(trajectories[0][i_step][6])
        
        trajectories = {
            'o'    : np.asarray(obs)[np.newaxis,:],
            'u'    : np.asarray(acts)[np.newaxis,:],
            'r'    : np.asarray(rews)[np.newaxis,:],
            'v'    : np.asarray(vals)[np.newaxis,:],
            'l'    : np.asarray(logs)[np.newaxis,:]
        }

    return trajectories, episode_reward, info['is_success'], frames

def run(model, experiment_args, train=True):

    total_time_start =  time.time()

    env, memory, noise, config, normalizer, _ = experiment_args
    
    N_EPISODES = config['n_episodes'] if train else config['n_episodes_test']
    N_CYCLES = config['n_cycles']
    N_ROLLOUTS = config['n_rollouts']
    N_BATCHES = config['n_batches']
    N_BD_BATCHES = config['n_bd_batches']
    N_TEST_ROLLOUTS = config['n_test_rollouts']
    BATCH_SIZE = config['batch_size']
    FT_RATE = config['ai_object_fine_tune_rate']
    
    episode_reward_all = []
    episode_success_all = []
    episode_reward_mean = []
    episode_success_mean = []
    critic_losses = []
    actor_losses = []
    dist_entropies = []
    backward_losses = []
        
    for i_episode in range(N_EPISODES):

        lr = config['plcy_lr']  - (config['plcy_lr']  * (i_episode / float(N_EPISODES)))
        for param_group in model.optims[0].param_groups:
            param_group['lr'] = lr
        
        episode_time_start = time.time()
        if train:
            for i_cycle in range(N_CYCLES):
                
                trajectories = None
                for i_rollout in range(N_ROLLOUTS):
                    # Initialize the environment and state
                    ai_object = 1 if np.random.rand() < config['ai_object_rate']  else 0
                    i_env = i_rollout % config['n_envs']
                    render = config['render'] > 0 and i_rollout==0
                    trajectory, _, _, _ = rollout(env, model, noise, i_env, normalizer, render=render, agent_id=0, ai_object=ai_object, rob_policy=config['rob_policy'])
                    if trajectories is None:
                        trajectories = trajectory.copy()
                    else:
                        for key in trajectories.keys():
                            trajectories[key] = np.concatenate((trajectories[key], trajectory.copy()[key]), axis=0)

                memory[0].obs.copy_(K.tensor(np.swapaxes(trajectories['o'], 0, 1), dtype=model.dtype, device=model.device))
                memory[0].actions.copy_(K.tensor(np.swapaxes(trajectories['u'], 0, 1), dtype=model.dtype, device=model.device))
                memory[0].rewards.copy_(K.tensor(np.swapaxes(trajectories['r'], 0, 1), dtype=model.dtype, device=model.device))
                memory[0].value_preds.copy_(K.tensor(np.swapaxes(trajectories['v'], 0, 1), dtype=model.dtype, device=model.device))
                memory[0].action_log_probs.copy_(K.tensor(np.swapaxes(trajectories['l'], 0, 1), dtype=model.dtype, device=model.device))

                memory[0].compute_returns(config['use_gae'], config['gamma'], config['gae_lambda'])

                model.to_cuda()
                memory[0].to(model.device)

                critic_loss, actor_loss, dist_entropy = model.update(memory[0])

                critic_losses.append(critic_loss)
                actor_losses.append(actor_loss)
                dist_entropies.append(dist_entropy)

                if i_cycle%FT_RATE==FT_RATE-1:

                    for i_rollout in range(N_ROLLOUTS):
                        i_env = i_rollout % config['n_envs']
                        render = config['render'] > 0 and i_rollout==0
                        trajectories, _, _, _ = rollout(env, model, noise, i_env, normalizer, render=render, agent_id=0, ai_object=True, rob_policy=config['rob_policy'])
                        memory[1].store_episode(trajectories.copy())
                        
                    for i_batch in range(N_BATCHES):  
                        model.to_cuda()
                        batch = memory[1].sample(BATCH_SIZE)
                        _, _ = model.update_object_parameters(batch, normalizer)

                    model.update_object_target()

                    for i_batch in range(N_BD_BATCHES):
                        batch = memory[1].sample(BATCH_SIZE)
                        _ = model.update_backward(batch, normalizer)  

            # <-- end loop: i_cycle
        plot_durations(np.asarray(critic_losses), np.asarray(critic_losses))
        plot_durations(np.asarray(actor_losses), np.asarray(actor_losses))
        plot_durations(np.asarray(dist_entropies), np.asarray(dist_entropies))
        plot_durations(np.asarray(backward_losses), np.asarray(backward_losses))

        episode_reward_cycle = []
        episode_succeess_cycle = []
        for i_rollout in range(N_TEST_ROLLOUTS):
            # Initialize the environment and state
            rollout_per_env = N_TEST_ROLLOUTS // config['n_envs']
            i_env = i_rollout // rollout_per_env
            render = config['render'] > 0 and i_episode % config['render'] == 0 and i_env == 0
            _, episode_reward, success, _ = rollout(env, model, (False,False), i_env, normalizer=normalizer, render=render, agent_id=0, ai_object=False, rob_policy=config['rob_policy'])
   
            episode_reward_cycle.append(episode_reward.item())
            episode_succeess_cycle.append(success)
        # <-- end loop: i_rollout 
            
        ### MONITORIRNG ###
        episode_reward_all.append(episode_reward_cycle)
        episode_success_all.append(episode_succeess_cycle)

        episode_reward_mean.append(np.mean(episode_reward_cycle))
        episode_success_mean.append(np.mean(episode_succeess_cycle))
        plot_durations(np.asarray(episode_reward_mean), np.asarray(episode_success_mean))
        
        if config['verbose'] > 0:
        # Printing out
            if (i_episode+1)%1 == 0:
                print("==> Episode {} of {}".format(i_episode + 1, N_EPISODES))
                print('  | Id exp: {}'.format(config['exp_id']))
                print('  | Exp description: {}'.format(config['exp_descr']))
                print('  | Env: {}'.format(config['env_id']))
                print('  | Process pid: {}'.format(config['process_pid']))
                print('  | Episode total reward: {}'.format(episode_reward))
                print('  | Running mean of total reward: {}'.format(episode_reward_mean[-1]))
                print('  | Success rate: {}'.format(episode_success_mean[-1]))
                #print('  | Running mean of total reward: {}'.format(running_mean(episode_reward_all)[-1]))
                print('  | Time episode: {}'.format(time.time()-episode_time_start))
                print('  | Time total: {}'.format(time.time()-total_time_start))
            
    # <-- end loop: i_episode
    if train:
        print('Training completed')
    else:
        print('Test completed')

    return episode_reward_all, episode_success_all

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display

plt.ion()

def plot_durations(p, r):
    plt.figure(2)
    plt.clf()
    plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Duration')
    plt.plot(p)
    plt.plot(r)

    plt.pause(0.001)  # pause a bit so that plots are updated
    if is_ipython:
        display.clear_output(wait=True)
        display.display(plt.gcf())<file_sep>import numpy as np
import scipy as sc
import time
import imageio

import torch as K
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F

import gym_wmgds as gym

from olbr.algorithms import DDPG, HDDPG
from olbr.experience import ReplayMemory, Transition
from olbr.exploration import Noise
from olbr.utils import Saver, Summarizer, get_noise_scale, get_params, running_mean
from olbr.agents.basic import Actor 
from olbr.agents.basic import Critic

import pdb

device = K.device("cuda" if K.cuda.is_available() else "cpu")
dtype = K.float32

def init(config):
    
    if config['resume'] != '':
        resume_path = config['resume']
        saver = Saver(config)
        config, start_episode, save_dict = saver.resume_ckpt()
        config['resume'] = resume_path
    else:
        start_episode = 0
        
    #hyperparameters
    ENV_NAME = config['env_id'] #'simple_spread'
    SEED = config['random_seed'] # 1

    GAMMA = config['gamma'] # 0.95
    TAU = config['tau'] # 0.01

    ACTOR_LR = config['plcy_lr'] # 0.01
    CRITIC_LR = config['crtc_lr'] # 0.01

    MEM_SIZE = config['buffer_length'] # 1000000 

    REGULARIZATION = config['regularization'] # True
    NORMALIZED_REWARDS = config['reward_normalization'] # True

    env = gym.make(ENV_NAME)
    env.seed(SEED)

    observation_space = env.observation_space.spaces['observation'].shape[0] + env.observation_space.spaces['desired_goal'].shape[0]
    #goal_space = gym.spaces.box.Box(low=-1, high=1, shape=(8,), dtype='float32')
    #goal_space = gym.spaces.box.Box(low=-2, high=2, shape=env.observation_space.spaces['desired_goal'].shape, dtype='float32')
    goal_space = gym.spaces.box.Box(low=-1, high=1, shape=env.observation_space.spaces['desired_goal'].shape, dtype='float32')
    action_space = env.action_space
    if env.action_space.low[0] == -1 and env.action_space.high[0] == 1:
        OUT_FUNC = K.tanh 
    elif env.action_space.low[0] == 0 and env.action_space.high[0] == 1:
        OUT_FUNC = K.sigmoid
    else:
        OUT_FUNC = K.sigmoid

    K.manual_seed(SEED)
    np.random.seed(SEED)
    
    MODEL = HDDPG

    if config['verbose'] > 1:
        # utils
        summaries = (Summarizer(config['dir_summary_train'], config['port'], config['resume']),
                    Summarizer(config['dir_summary_test'], config['port'], config['resume']))
        saver = Saver(config)
    else:
        summaries = None
        saver = None

    #exploration initialization
    noise = (Noise(action_space.shape[0], sigma=0.05, eps=0.2),
             Noise(goal_space.shape[0], sigma=0.05, eps=0.2))
    #noise = OUNoise(action_space.shape[0])

    #model initialization
    optimizer = (optim.Adam, (ACTOR_LR, CRITIC_LR)) # optimiser func, (actor_lr, critic_lr)
    loss_func = F.mse_loss
    model = MODEL(observation_space, action_space, goal_space, optimizer, 
                  Actor, Critic, loss_func, GAMMA, TAU, out_func=OUT_FUNC,
                  discrete=False, regularization=REGULARIZATION, normalized_rewards=NORMALIZED_REWARDS)
    
    if config['resume'] != '':
        for i, param in enumerate(save_dict['model_params']):
            model.entities[i].load_state_dict(param)
    
    #memory initilization
    memory = [ReplayMemory(MEM_SIZE), ReplayMemory(MEM_SIZE)]

    experiment_args = (env, memory, noise, config, summaries, saver, start_episode)
          
    return model, experiment_args

def rollout(env, model, noise, mask, ratio, mapping, C=25, render=False, sparse_rewards=True):
    trajectory = [[], []]

    # monitoring variables
    episode_reward = [0, 0]
    frames = []
    
    state = env.reset()
    success = 0
    
    for i_step in range(env._max_episode_steps):

        model.to_cpu()

        obs = K.tensor(state['observation'], dtype=K.float32).unsqueeze(0)
        goal = K.tensor(state['desired_goal'], dtype=K.float32).unsqueeze(0)
        obs = K.cat([obs, goal], dim=-1)

        obs_f = obs*mask
        obs_s = obs#*(1-mask)

        if i_step % C  == 0:
            state_init = state.copy()
            #goal_f = model.select_goal(obs_s, noise[1])
            #pdb.set_trace()
            goal_f = goal
            #goal_f = goal - obs[:,0:3]
            #goal_f *= (ratio[0:3]*C*2)
            #goal_f += obs_s[:,0:3]
            #pdb.set_trace()
            #goal_f *= (ratio[K.tensor(mask, dtype=K.uint8)]*C)
            #goal_f += obs_s[:,K.tensor(mask, dtype=K.uint8)]            
            cumm_reward = K.zeros((1,1), dtype=K.float32)

        action = model.select_action(K.cat([obs_f, goal_f], dim=-1), noise[0])

        next_state, reward, done, info = env.step(action.squeeze(0).numpy())

        next_obs = K.tensor(next_state['observation'], dtype=K.float32).unsqueeze(0)
        #next_goal = K.tensor(next_state['desired_goal'], dtype=K.float32).unsqueeze(0)
        next_obs = K.cat([next_obs, goal], dim=-1)

        next_obs_f = next_obs*mask
        #next_obs_s = next_obs*(1-mask)

        reward = K.tensor(reward, dtype=K.float32).view(1,1)
        cumm_reward += reward
        #old_distance = F.mse_loss(goal_f, mapping(obs_f, mask)).view(1,1)
        #new_distance = F.mse_loss(goal_f, mapping(next_obs_f, mask)).view(1,1)
        #old_distance = goal_f - mapping(obs_f, mask)
        #new_distance = goal_f - mapping(next_obs_f, mask)
        #old_distance /= (ratio[K.tensor(mask, dtype=K.uint8)]*C)
        #new_distance /= (ratio[K.tensor(mask, dtype=K.uint8)]*C)

        #old_distance = K.tensor(np.linalg.norm(old_distance, axis=-1), dtype=K.float32).view(1,1)
        #new_distance = K.tensor(np.linalg.norm(new_distance, axis=-1), dtype=K.float32).view(1,1)
        distance = K.tensor(np.linalg.norm(mapping(goal_f) - mapping(next_obs_f), axis=-1), dtype=K.float32).view(1,1)
        #distance = K.tensor(np.linalg.norm(mapping(goal_f) - 0, axis=-1), dtype=K.float32).view(1,1)
        success = distance < 0.05
        if sparse_rewards:
            intr_reward = - K.tensor(distance > 0.05, dtype=K.float32)
        else:
            intr_reward = old_distance - new_distance
        #K.sqrt(K.pow((goal_f - mapping(next_obs_f)), 2).sum(-1))
        
        #print(K.sqrt(K.pow((goal_f - mapping(next_obs_f)), 2).sum(-1)))
        #print(F.mse_loss(goal_f, mapping(next_obs_f)).view(1,1))

        # for monitoring
        episode_reward[0] = episode_reward[0] + intr_reward
        episode_reward[1] = episode_reward[1] + reward

        trajectory[0].append((state, action, intr_reward, next_state, goal_f, done))
        if (i_step+1) % C  == 0:
            trajectory[1].append((state_init, None, cumm_reward, next_state, goal_f, done))

        # Move to the next state
        state = next_state

        # Record frames
        if render:
            frames.append(env.render(mode='rgb_array')[0])

    return trajectory, episode_reward, (success, info['is_success']), frames, 
  
# def get_of_mask(env, nb_episodes=400, nb_steps=50, threshold=0.08):
#     states = []
#     actions = []
#     for _ in range(nb_episodes):
#         env.reset()
#         for _ in range(nb_steps):    
#             action = env.action_space.sample()
#             state, _, _, _ = env.step(action)
#             observation = state['observation']
#             desired_goal = state['desired_goal']
#             state = np.concatenate([observation, desired_goal])
#             states.append(state)
#             actions.append(action)

#     states = np.asarray(states)
#     actions = np.asarray(actions)

#     corr = np.zeros((states.shape[1], actions.shape[1]))
#     for i in range(states.shape[1]):
#         for j in range(actions.shape[1]):    
#             corr[i,j] = np.corrcoef(states[:,i], actions[:,j])[0,1]
    
#     mask = (K.tensor(abs(corr))>threshold).any(dim=1)
    
#     return mask

def get_of_mask(env, nb_episodes=400, nb_steps=50, threshold=0.08):
    states = []
    actions = []
    for _ in range(nb_episodes):
        old_state = env.reset()
        for _ in range(nb_steps):    
            action = env.action_space.sample()
            new_state, _, _, _ = env.step(action)
            observation_diff = new_state['observation'] - old_state['observation']
            desired_goal_diff = new_state['desired_goal'] - old_state['desired_goal'] 
            state = np.concatenate([observation_diff, desired_goal_diff])
            states.append(state)
            actions.append(action)
            old_state = new_state.copy()

    states = np.asarray(states)
    actions = np.asarray(actions)

    corr = np.zeros((states.shape[1], actions.shape[1]))
    for i in range(states.shape[1]):
        for j in range(actions.shape[1]):    
            corr[i,j] = np.corrcoef(states[:,i], actions[:,j])[0,1]
            
    ratio = np.zeros((states.shape[1], actions.shape[1]))
    for i in range(states.shape[1]):
        for j in range(actions.shape[1]):    
            ratio[i,j] = np.median(states[:,i]/actions[:,j])
    
    mask = (K.tensor(abs(corr))>threshold).any(dim=1)
    
    return K.tensor(mask, dtype=K.float32), K.max(K.tensor(ratio, dtype=K.float32), dim=-1)[0]

def mapping(x, mask=None):
    if mask is None:
        return x[:,0:3]
    else:
        return x[:, K.tensor(mask, dtype=K.uint8)]

def run(model, experiment_args, train=True):

    total_time_start =  time.time()

    env, memory, noise, config, summaries, saver, start_episode = experiment_args
    
    start_episode = start_episode if train else 0
    NUM_EPISODES = config['n_episodes'] if train else config['n_episodes_test'] 
    EPISODE_LENGTH = env._max_episode_steps #config['episode_length'] if train else config['episode_length_test'] 
    
    episode_intr_reward_all = []
    episode_reward_all = []
    episode_intr_success_all = []
    episode_success_all = []

    mask, ratio = get_of_mask(env, nb_episodes=40)

    hier_C = 1
    sparse_rewards = True
        
    for i_episode in range(start_episode, NUM_EPISODES*5):
        
        episode_time_start = time.time()
        #noise.scale = get_noise_scale(i_episode, config)
        if train:
            for i_cycle in range(10):
            
                trajectories = []
                for i_rollout in range(16):
                    # Initialize the environment and state
                    trajectory, _, _, _ = rollout(env, model, noise, mask, ratio, mapping, hier_C, render=0, sparse_rewards=sparse_rewards)
                    trajectories.append(trajectory)

                for trajectory in trajectories:

                    for i_step in range(len(trajectory[0])):
                        state, action, reward, next_state, goal_f, done = trajectory[0][i_step]
                        
                        obs = K.tensor(state['observation'], dtype=K.float32).unsqueeze(0)
                        goal = K.tensor(state['desired_goal'], dtype=K.float32).unsqueeze(0)
                        obs = K.cat([obs, goal], dim=-1)
                        obs_f = obs*mask

                        next_obs = K.tensor(next_state['observation'], dtype=K.float32).unsqueeze(0)
                        next_obs = K.cat([next_obs, goal], dim=-1)
                        next_obs_f = next_obs*mask
                        next_achieved_f = mapping(next_obs_f)
                        
                        # regular sample
                        obs_goal_f = K.cat([obs_f, goal_f], dim=-1)
                        if done:
                            next_obs_goal_f = None
                        else:
                            next_obs_goal_f = K.cat([next_obs_f, goal_f], dim=-1)
                        
                        memory[0].push(obs_goal_f, action, next_obs_goal_f, reward)

                        # HER sample 
                        for _ in range(4):
                            future = np.random.randint(i_step, len(trajectory[0]))
                            _, _, _, next_state, _, _ = trajectory[0][future]

                            aux_goal_f = mapping(K.cat([K.tensor(next_state['observation'], dtype=K.float32).unsqueeze(0), goal], dim=-1)*mask)

                            obs_goal_f = K.cat([obs_f, aux_goal_f], dim=-1)#*(1-mask)
                            if done:
                                next_obs_goal_f = None
                            else:
                                next_obs_goal_f = K.cat([next_obs_f, aux_goal_f], dim=-1)#*(1-mask)
                                
                            reward = env.compute_reward(next_achieved_f, aux_goal_f, None)
                            reward = K.tensor(reward, dtype=dtype).view(1,1)
                            
                            memory[0].push(obs_goal_f, action, next_obs_goal_f, reward)

                    for i_step in range(len(trajectory[1])):
                        state, _, reward, next_state, goal_f, done = trajectory[1][i_step]
                        
                        obs = K.tensor(state['observation'], dtype=K.float32).unsqueeze(0)
                        goal = K.tensor(state['desired_goal'], dtype=K.float32).unsqueeze(0)

                        next_obs = K.tensor(next_state['observation'], dtype=K.float32).unsqueeze(0)
                        next_achieved = K.tensor(next_state['achieved_goal'], dtype=K.float32).unsqueeze(0)
                        
                        # regular sample
                        obs_goal_s = K.cat([obs, goal], dim=-1)#*(1-mask)
                        if done:
                            next_obs_goal_s = None
                        else:
                            next_obs_goal_s = K.cat([next_obs, goal], dim=-1)#*(1-mask)
                        
                        memory[1].push(obs_goal_s, goal_f, next_obs_goal_s, reward)
                        # HER sample 
                        for _ in range(4):
                            future = np.random.randint(i_step, len(trajectory[1]))
                            _, _, _, next_state, _, _ = trajectory[1][future]
                            aux_goal = K.tensor(next_state['achieved_goal'], dtype=K.float32).unsqueeze(0)

                            obs_goal_s = K.cat([obs, aux_goal], dim=-1)#*(1-mask)
                            if done:
                                next_obs_goal_s = None
                            else:
                                next_obs_goal_s = K.cat([next_obs, aux_goal], dim=-1)#*(1-mask)
                                
                            reward = env.compute_reward(next_achieved, aux_goal, None)
                            #reward = K.tensor(reward, dtype=dtype).view(1,1)*(hier_C/2) - hier_C/2
                            reward = K.tensor(reward, dtype=dtype).view(1,1)
                            
                            memory[1].push(obs_goal_s, goal_f, next_obs_goal_s, reward)

                    # <-- end loop: i_step
                # <-- end loop: i_rollout 

                for _ in range(40):
                    if len(memory[0]) > config['batch_size']-1 and len(memory[1]) > config['batch_size']-1:  
                        batch = []
                        model.to_cuda()  
                        batch.append(Transition(*zip(*memory[0].sample(config['batch_size']))))
                        batch.append(Transition(*zip(*memory[1].sample(config['batch_size']))))
                        critic_loss, actor_loss = model.update_parameters(batch)

            # <-- end loop: i_cycle

        episode_intr_reward_cycle = []
        episode_reward_cycle = []
        episode_intr_success_cycle = []
        episode_succeess_cycle = []
        for i_rollout in range(10):
            # Initialize the environment and state
            render = config['render'] > 0 and i_episode % config['render'] == 0
            _, episode_reward, success, frames = rollout(env, model, noise=(False, False), mask=mask, ratio=ratio, mapping=mapping, C=hier_C, render=render, sparse_rewards=sparse_rewards)
                
            # Save gif
            dir_monitor = config['dir_monitor_train'] if train else config['dir_monitor_test']
            if config['render'] > 0 and i_episode % config['render'] == 0:
                imageio.mimsave('{}/{}.gif'.format(dir_monitor, i_episode), frames)

            
            episode_intr_reward_cycle.append(episode_reward[0])
            episode_reward_cycle.append(episode_reward[1])
            episode_intr_success_cycle.append(success[0])
            episode_succeess_cycle.append(success[1])
        # <-- end loop: i_rollout 
            
        ### MONITORIRNG ###
        episode_intr_reward_all.append(np.mean(episode_intr_reward_cycle))
        episode_reward_all.append(np.mean(episode_reward_cycle))
        episode_intr_success_all.append(np.mean(episode_intr_success_cycle))
        episode_success_all.append(np.mean(episode_succeess_cycle))
        
        if config['verbose'] > 0:
        # Printing out
            if (i_episode+1)%1 == 0:
                print("==> Episode {} of {}".format(i_episode + 1, NUM_EPISODES))
                print('  | Id exp: {}'.format(config['exp_id']))
                print('  | Exp description: {}'.format(config['exp_descr']))
                print('  | Env: {}'.format(config['env_id']))
                print('  | Process pid: {}'.format(config['process_pid']))
                print('  | Tensorboard port: {}'.format(config['port']))
                print('  | Episode total reward: {}'.format(episode_reward))
                print('  | Running mean of total reward: {}'.format(episode_reward_all[-1]))
                print('  | Running mean of total intrinsic reward: {}'.format(episode_intr_reward_all[-1]))
                print('  | Success rate: {}'.format(episode_success_all[-1]))
                print('  | Intrinsic success rate: {}'.format(episode_intr_success_all[-1]))
                #print('  | Running mean of total reward: {}'.format(running_mean(episode_reward_all)[-1]))
                print('  | Time episode: {}'.format(time.time()-episode_time_start))
                print('  | Time total: {}'.format(time.time()-total_time_start))
                        
        if config['verbose'] > 0:    
            ep_save = i_episode+1 if (i_episode == NUM_EPISODES-1) else None  
            is_best_save = None
            is_best_avg_save = None   
                
            if (not train) or ((np.asarray([ep_save, is_best_save, is_best_avg_save]) == None).sum() == 3):
                to_save = False
            else:
                model.to_cpu()
                saver.save_checkpoint(save_dict   = {'model_params': [entity.state_dict() for entity in model.entities]},
                                      episode     = ep_save,
                                      is_best     = is_best_save,
                                      is_best_avg = is_best_avg_save
                                      )
                to_save = True
    
            #if (i_episode+1)%100 == 0:
            #    summary = summaries[0] if train else summaries[1]
            #    summary.update_log(i_episode, 
            #                    episode_reward, 
            #                    list(episode_reward.reshape(-1,)), 
            #                    critic_loss        = critic_loss, 
            #                    actor_loss         = actor_loss,
            #                    to_save            = to_save
            #                    )
        


            
    # <-- end loop: i_episode
    if train:
        print('Training completed')
    else:
        print('Test completed')
    
    return episode_reward_all


if __name__ == '__main__':

    monitor_macddpg_p2 = []
    monitor_macddpg_p2_test = []
    for i in range(0,5):
        config = get_params(args=['--exp_id','MACDDPG_P2_120K_'+ str(i+1), 
                                '--random_seed', str(i+1), 
                                '--agent_alg', 'MACDDPG',
                                '--protocol_type', str(2),
                                '--n_episodes', '120000',
                                '--verbose', '2',
                                ]
                        )
        model, experiment_args = init(config)

        env, memory, ounoise, config, summaries, saver, start_episode = experiment_args

        tic = time.time()
        monitor = run(model, experiment_args, train=True)
        monitor_test = run(model, experiment_args, train=False)

        toc = time.time()

        env.close()
        for summary in summaries:
            summary.close()
            
        monitor_macddpg_p2.append(monitor)
        monitor_macddpg_p2_test.append(monitor_test)
        
        np.save('./monitor_macddpg_p2.npy', monitor_macddpg_p2)
        np.save('./monitor_macddpg_p2_test.npy', monitor_macddpg_p2_test)
        
        print(toc-tic)