コード例 #1
0
    def __init__(self, args):
        if args.seed is not None:
            torch.manual_seed(args.seed)
            np.random.seed(args.seed)
            random.seed(args.seed)

        self.env_name = args.environment
        self.env_setting = get_env_setting(self.env_name)
        self.solved_reward = self.env_setting["solved_reward"]
        self.update_timestep = self.env_setting["update_timestep"]
        self.device = torch.device(
            "cuda:0" if torch.cuda.is_available() else "cpu")
        self.env = gym.make(args.environment)
        self.alg = PPO(args, self.device)
        self.log_interval = 5  # print avg reward in the interval
        self.max_episodes = 100000
        self.render = False
コード例 #2
0
def main():
    args = get_config()

    # seed
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)

    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(args.n_training_threads)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)

    # path
    model_dir = Path('./results') / args.env_name / args.algorithm_name
    if not model_dir.exists():
        curr_run = 'run1'
    else:
        exst_run_nums = [
            int(str(folder.name).split('run')[1])
            for folder in model_dir.iterdir()
            if str(folder.name).startswith('run')
        ]
        if len(exst_run_nums) == 0:
            curr_run = 'run1'
        else:
            curr_run = 'run%i' % (max(exst_run_nums) + 1)

    run_dir = model_dir / curr_run
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir))

    # env
    envs = make_parallel_env(args)
    #Policy network
    actor_critic = []
    if args.share_policy:
        ac = Policy(envs.observation_space[0],
                    envs.action_space[0],
                    num_agents=args.num_agents,
                    base_kwargs={
                        'lstm': args.lstm,
                        'naive_recurrent': args.naive_recurrent_policy,
                        'recurrent': args.recurrent_policy,
                        'hidden_size': args.hidden_size
                    })
        ac.to(device)
        for agent_id in range(args.num_agents):
            actor_critic.append(ac)
    else:
        for agent_id in range(args.num_agents):
            ac = Policy(envs.observation_space[0],
                        envs.action_space[0],
                        num_agents=args.num_agents,
                        base_kwargs={
                            'naive_recurrent': args.naive_recurrent_policy,
                            'recurrent': args.recurrent_policy,
                            'hidden_size': args.hidden_size
                        })
            ac.to(device)
            actor_critic.append(ac)

    agents = []
    rollouts = []
    for agent_id in range(args.num_agents):
        # algorithm
        agent = PPO(actor_critic[agent_id],
                    agent_id,
                    args.clip_param,
                    args.ppo_epoch,
                    args.num_mini_batch,
                    args.data_chunk_length,
                    args.value_loss_coef,
                    args.entropy_coef,
                    logger,
                    lr=args.lr,
                    eps=args.eps,
                    max_grad_norm=args.max_grad_norm,
                    use_clipped_value_loss=args.use_clipped_value_loss)
        #replay buffer
        ro = RolloutStorage(args.num_agents, agent_id, args.episode_length,
                            args.n_rollout_threads,
                            envs.observation_space[agent_id],
                            envs.action_space[agent_id],
                            actor_critic[agent_id].recurrent_hidden_state_size)
        agents.append(agent)
        rollouts.append(ro)

    # reset env
    obs = envs.reset()
    # rollout
    for i in range(args.num_agents):
        rollouts[i].share_obs[0].copy_(
            torch.tensor(obs.reshape(args.n_rollout_threads, -1)))
        rollouts[i].obs[0].copy_(torch.tensor(obs[:, i, :]))
        rollouts[i].recurrent_hidden_states.zero_()
        rollouts[i].recurrent_hidden_states_critic.zero_()
        rollouts[i].recurrent_c_states.zero_()
        rollouts[i].recurrent_c_states_critic.zero_()
        rollouts[i].to(device)

    # run
    coop_num = []
    defect_num = []
    coopdefect_num = []
    defectcoop_num = []
    gore1_num = []
    gore2_num = []
    gore3_num = []
    hare1_num = []
    hare2_num = []
    hare3_num = []
    collective_return = []
    apple_consumption = []
    waste_cleared = []
    sustainability = []
    fire = []

    start = time.time()
    episodes = int(
        args.num_env_steps) // args.episode_length // args.n_rollout_threads
    all_episode = 0

    for episode in range(episodes):

        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            for i in range(args.num_agents):
                update_linear_schedule(agents[i].optimizer, episode, episodes,
                                       args.lr)

        for step in range(args.episode_length):
            # Sample actions
            values = []
            actions = []
            action_log_probs = []
            recurrent_hidden_statess = []
            recurrent_hidden_statess_critic = []
            recurrent_c_statess = []
            recurrent_c_statess_critic = []

            with torch.no_grad():
                for i in range(args.num_agents):
                    value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic,\
                        recurrent_c_states, recurrent_c_states_critic =\
                            actor_critic[i].act(rollouts[i].share_obs[step],
                                                        rollouts[i].obs[step],
                                                        rollouts[i].recurrent_hidden_states[step],
                                                        rollouts[i].recurrent_hidden_states_critic[step],
                                                        rollouts[i].recurrent_c_states[step],
                                                        rollouts[i].recurrent_c_states_critic[step],
                                                        rollouts[i].masks[step])
                    values.append(value)
                    actions.append(action)
                    action_log_probs.append(action_log_prob)
                    recurrent_hidden_statess.append(recurrent_hidden_states)
                    recurrent_hidden_statess_critic.append(
                        recurrent_hidden_states_critic)
                    recurrent_c_statess.append(recurrent_c_states)
                    recurrent_c_statess_critic.append(
                        recurrent_c_states_critic)

            # rearrange action
            actions_env = []
            for i in range(args.n_rollout_threads):
                one_hot_action_env = []
                for k in range(args.num_agents):
                    one_hot_action = np.zeros(envs.action_space[0].n)
                    one_hot_action[actions[k][i]] = 1
                    one_hot_action_env.append(one_hot_action)
                actions_env.append(one_hot_action_env)

            # Obser reward and next obs
            obs, reward, done, infos = envs.step(actions_env)

            # If done then clean the history of observations.
            # insert data in buffer
            masks = []
            bad_masks = []
            for i in range(args.num_agents):
                mask = []
                bad_mask = []
                for done_ in done:
                    if done_[i]:
                        mask.append([0.0])
                        bad_mask.append([1.0])
                    else:
                        mask.append([1.0])
                        bad_mask.append([1.0])
                masks.append(torch.FloatTensor(mask))
                bad_masks.append(torch.FloatTensor(bad_mask))

            for i in range(args.num_agents):
                rollouts[i].insert(
                    torch.tensor(obs.reshape(args.n_rollout_threads, -1)),
                    torch.tensor(obs[:, i, :]), recurrent_hidden_statess[i],
                    recurrent_hidden_statess_critic[i], recurrent_c_statess[i],
                    recurrent_c_statess_critic[i], actions[i],
                    action_log_probs[i], values[i],
                    torch.tensor(reward[:,
                                        i].reshape(-1,
                                                   1)), masks[i], bad_masks[i])

        with torch.no_grad():
            next_values = []
            for i in range(args.num_agents):
                next_value = actor_critic[i].get_value(
                    rollouts[i].share_obs[-1], rollouts[i].obs[-1],
                    rollouts[i].recurrent_hidden_states[-1],
                    rollouts[i].recurrent_hidden_states_critic[-1],
                    rollouts[i].recurrent_c_states[-1],
                    rollouts[i].recurrent_c_states_critic[-1],
                    rollouts[i].masks[-1]).detach()
                next_values.append(next_value)

        for i in range(args.num_agents):
            rollouts[i].compute_returns(next_values[i], args.use_gae,
                                        args.gamma, args.gae_lambda,
                                        args.use_proper_time_limits)

        # update the network
        value_losses = []
        action_losses = []
        dist_entropies = []
        for i in range(args.num_agents):
            value_loss, action_loss, dist_entropy = agents[i].update(
                rollouts[i])
            value_losses.append(value_loss)
            action_losses.append(action_loss)
            dist_entropies.append(dist_entropy)

        if args.env_name == "StagHunt":
            for info in infos:
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'defect&defect_num' in info.keys():
                    defect_num.append(info['defect&defect_num'])
                if 'coop&defect_num' in info.keys():
                    coopdefect_num.append(info['coop&defect_num'])
                if 'defect&coop_num' in info.keys():
                    defectcoop_num.append(info['defect&coop_num'])

            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'defect&defect_num_per_episode',
                    {'defect&defect_num_per_episode': defect_num[all_episode]},
                    all_episode)
                logger.add_scalars('coop&defect_num_per_episode', {
                    'coop&defect_num_per_episode':
                    coopdefect_num[all_episode]
                }, all_episode)
                logger.add_scalars('defect&coop_num_per_episode', {
                    'defect&coop_num_per_episode':
                    defectcoop_num[all_episode]
                }, all_episode)
                all_episode += 1
        elif args.env_name == "StagHuntGW":
            for info in infos:
                if 'collective_return' in info.keys():
                    collective_return.append(info['collective_return'])
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'gore1_num' in info.keys():
                    gore1_num.append(info['gore1_num'])
                if 'gore2_num' in info.keys():
                    gore2_num.append(info['gore2_num'])
                if 'hare1_num' in info.keys():
                    hare1_num.append(info['hare1_num'])
                if 'hare2_num' in info.keys():
                    hare2_num.append(info['hare2_num'])
            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'collective_return',
                    {'collective_return': collective_return[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore1_num_per_episode',
                    {'gore1_num_per_episode': gore1_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore2_num_per_episode',
                    {'gore2_num_per_episode': gore2_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'hare1_num_per_episode',
                    {'hare1_num_per_episode': hare1_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'hare2_num_per_episode',
                    {'hare2_num_per_episode': hare2_num[all_episode]},
                    all_episode)
                all_episode += 1
        elif args.env_name == "EscalationGW":
            for info in infos:
                if 'collective_return' in info.keys():
                    collective_return.append(info['collective_return'])
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'collective_return',
                    {'collective_return': collective_return[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[all_episode]},
                    all_episode)
                all_episode += 1
        elif args.env_name == "multi_StagHuntGW":
            for info in infos:
                if 'collective_return' in info.keys():
                    collective_return.append(info['collective_return'])
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'gore0_num' in info.keys():
                    gore1_num.append(info['gore0_num'])
                if 'gore1_num' in info.keys():
                    gore2_num.append(info['gore1_num'])
                if 'gore2_num' in info.keys():
                    gore3_num.append(info['gore2_num'])
                if 'hare0_num' in info.keys():
                    hare1_num.append(info['hare0_num'])
                if 'hare1_num' in info.keys():
                    hare2_num.append(info['hare1_num'])
                if 'hare2_num' in info.keys():
                    hare3_num.append(info['hare2_num'])
            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'collective_return',
                    {'collective_return': collective_return[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore1_num_per_episode',
                    {'gore1_num_per_episode': gore1_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore2_num_per_episode',
                    {'gore2_num_per_episode': gore2_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore3_num_per_episode',
                    {'gore3_num_per_episode': gore3_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'hare1_num_per_episode',
                    {'hare1_num_per_episode': hare1_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'hare2_num_per_episode',
                    {'hare2_num_per_episode': hare2_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'hare3_num_per_episode',
                    {'hare3_num_per_episode': hare3_num[all_episode]},
                    all_episode)
                all_episode += 1

        # clean the buffer and reset
        obs = envs.reset()
        for i in range(args.num_agents):
            rollouts[i].share_obs[0].copy_(
                torch.tensor(obs.reshape(args.n_rollout_threads, -1)))
            rollouts[i].obs[0].copy_(torch.tensor(obs[:, i, :]))
            rollouts[i].recurrent_hidden_states.zero_()
            rollouts[i].recurrent_hidden_states_critic.zero_()
            rollouts[i].recurrent_c_states.zero_()
            rollouts[i].recurrent_c_states_critic.zero_()
            rollouts[i].masks[0].copy_(torch.ones(args.n_rollout_threads, 1))
            rollouts[i].bad_masks[0].copy_(
                torch.ones(args.n_rollout_threads, 1))
            rollouts[i].to(device)

        for i in range(args.num_agents):
            # save for every interval-th episode or for the last epoch
            if (episode % args.save_interval == 0 or episode == episodes - 1):
                torch.save({'model': actor_critic[i]},
                           str(save_dir) + "/agent%i_model" % i + ".pt")

        # log information
        if episode % args.log_interval == 0:
            total_num_steps = (
                episode + 1) * args.episode_length * args.n_rollout_threads
            end = time.time()
            print(
                "\n Updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(episode, episodes, total_num_steps, args.num_env_steps,
                        int(total_num_steps / (end - start))))
            for i in range(args.num_agents):
                print("value loss of agent%i: " % i + str(value_losses[i]))
    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()

    ###----------------------------------------------------------###
    ###----------------------------------------------------------###
    ###----------------------------------------------------------###
    if args.eval:
        eval_dir = run_dir / 'eval'
        log_dir = eval_dir / 'logs'
        os.makedirs(str(log_dir))
        logger = SummaryWriter(str(log_dir))

        # eval best policy
        eval_rewards = []
        # env
        if args.env_name == "StagHunt":
            assert args.num_agents == 2, (
                "only 2 agents is supported, check the config.py.")
            env = MGEnv(args)
        elif args.env_name == "StagHuntGW" or args.env_name == "EscalationGW":
            assert args.num_agents == 2, (
                "only 2 agent is supported in single navigation, check the config.py."
            )
            env = GridWorldEnv(args)
        elif args.env_name == "multi_StagHuntGW":
            env = multi_GridWorldEnv(args)
        else:
            print("Can not support the " + args.env_name + "environment.")
            raise NotImplementedError

        #Policy network
        coop_num = []
        defect_num = []
        coopdefect_num = []
        defectcoop_num = []
        gore1_num = []
        gore2_num = []
        gore3_num = []
        hare1_num = []
        hare2_num = []
        hare3_num = []
        collective_return = []
        apple_consumption = []
        waste_cleared = []
        sustainability = []
        fire = []

        for episode in range(args.eval_episodes):
            print("Episode %i of %i" % (episode, args.eval_episodes))
            state = env.reset()
            state = np.array([state])

            share_obs = []
            obs = []
            recurrent_hidden_statess = []
            recurrent_hidden_statess_critic = []
            recurrent_c_statess = []
            recurrent_c_statess_critic = []
            masks = []
            policy_reward = 0

            # rollout
            for i in range(args.num_agents):
                share_obs.append(
                    (torch.tensor(state.reshape(1, -1),
                                  dtype=torch.float32)).to(device))
                obs.append((torch.tensor(state[:, i, :],
                                         dtype=torch.float32)).to(device))
                recurrent_hidden_statess.append(
                    torch.zeros(
                        1, actor_critic[i].recurrent_hidden_state_size).to(
                            device))
                recurrent_hidden_statess_critic.append(
                    torch.zeros(
                        1, actor_critic[i].recurrent_hidden_state_size).to(
                            device))
                recurrent_c_statess.append(
                    torch.zeros(
                        1, actor_critic[i].recurrent_hidden_state_size).to(
                            device))
                recurrent_c_statess_critic.append(
                    torch.zeros(
                        1, actor_critic[i].recurrent_hidden_state_size).to(
                            device))
                masks.append(torch.ones(1, 1).to(device))

            for step in range(args.episode_length):
                print("step %i of %i" % (step, args.episode_length))
                # Sample actions
                one_hot_actions = []
                for i in range(args.num_agents):
                    one_hot_action = np.zeros(env.action_space[0].n)
                    with torch.no_grad():
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic, recurrent_c_states, recurrent_c_states_critic = actor_critic[
                            i].act(share_obs[i], obs[i],
                                   recurrent_hidden_statess[i],
                                   recurrent_hidden_statess_critic[i],
                                   recurrent_c_statess[i],
                                   recurrent_c_statess_critic[i], masks[i])
                    recurrent_hidden_statess[i].copy_(recurrent_hidden_states)
                    recurrent_hidden_statess_critic[i].copy_(
                        recurrent_hidden_states_critic)
                    recurrent_c_statess[i].copy_(recurrent_c_states)
                    recurrent_c_statess_critic[i].copy_(
                        recurrent_c_states_critic)
                    one_hot_action[action] = 1
                    one_hot_actions.append(one_hot_action)

                # Obser reward and next obs
                state, reward, done, infos = env.step(one_hot_actions)

                for i in range(args.num_agents):
                    print("Reward of agent%i: " % i + str(reward[i]))
                    policy_reward += reward[i]

                if all(done):
                    break

                state = np.array([state])

                for i in range(args.num_agents):
                    if len(env.observation_space[0]) == 1:
                        share_obs[i].copy_(
                            torch.tensor(state.reshape(1, -1),
                                         dtype=torch.float32))
                        obs[i].copy_(
                            torch.tensor(state[:, i, :], dtype=torch.float32))
                    elif len(env.observation_space[0]) == 3:
                        share_obs[i].copy_(
                            torch.tensor(state.reshape(
                                1, -1, env.observation_space[0][1],
                                env.observation_space[0][2]),
                                         dtype=torch.float32))
                        obs[i].copy_(
                            torch.tensor(state[:, i, :, :, :],
                                         dtype=torch.float32))

            eval_rewards.append(policy_reward)

            if args.env_name == "StagHunt":
                if 'coop&coop_num' in infos.keys():
                    coop_num.append(infos['coop&coop_num'])
                if 'defect&defect_num' in infos.keys():
                    defect_num.append(infos['defect&defect_num'])
                if 'coop&defect_num' in infos.keys():
                    coopdefect_num.append(infos['coop&defect_num'])
                if 'defect&coop_num' in infos.keys():
                    defectcoop_num.append(infos['defect&coop_num'])

                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[episode]}, episode)
                logger.add_scalars(
                    'defect&defect_num_per_episode',
                    {'defect&defect_num_per_episode': defect_num[episode]},
                    episode)
                logger.add_scalars(
                    'coop&defect_num_per_episode',
                    {'coop&defect_num_per_episode': coopdefect_num[episode]},
                    episode)
                logger.add_scalars(
                    'defect&coop_num_per_episode',
                    {'defect&coop_num_per_episode': defectcoop_num[episode]},
                    episode)

            elif args.env_name == "StagHuntGW":
                if 'collective_return' in infos.keys():
                    collective_return.append(infos['collective_return'])
                    logger.add_scalars(
                        'collective_return',
                        {'collective_return': collective_return[episode]},
                        episode)
                if 'coop&coop_num' in infos.keys():
                    coop_num.append(infos['coop&coop_num'])
                    logger.add_scalars(
                        'coop&coop_num_per_episode',
                        {'coop&coop_num_per_episode': coop_num[episode]},
                        episode)
                if 'gore1_num' in infos.keys():
                    gore1_num.append(infos['gore1_num'])
                    logger.add_scalars(
                        'gore1_num_per_episode',
                        {'gore1_num_per_episode': gore1_num[episode]}, episode)
                if 'gore2_num' in infos.keys():
                    gore2_num.append(infos['gore2_num'])
                    logger.add_scalars(
                        'gore2_num_per_episode',
                        {'gore2_num_per_episode': gore2_num[episode]}, episode)
                if 'hare1_num' in infos.keys():
                    hare1_num.append(infos['hare1_num'])
                    logger.add_scalars(
                        'hare1_num_per_episode',
                        {'hare1_num_per_episode': hare1_num[episode]}, episode)
                if 'hare2_num' in infos.keys():
                    hare2_num.append(infos['hare2_num'])
                    logger.add_scalars(
                        'hare2_num_per_episode',
                        {'hare2_num_per_episode': hare2_num[episode]}, episode)
            elif args.env_name == "EscalationGW":
                if 'collective_return' in infos.keys():
                    collective_return.append(infos['collective_return'])
                    logger.add_scalars(
                        'collective_return',
                        {'collective_return': collective_return[episode]},
                        episode)
                if 'coop&coop_num' in infos.keys():
                    coop_num.append(infos['coop&coop_num'])
                    logger.add_scalars(
                        'coop&coop_num_per_episode',
                        {'coop&coop_num_per_episode': coop_num[episode]},
                        episode)
            elif args.env_name == "multi_StagHuntGW":
                if 'collective_return' in infos.keys():
                    collective_return.append(infos['collective_return'])
                    logger.add_scalars(
                        'collective_return',
                        {'collective_return': collective_return[episode]},
                        episode)
                if 'coop&coop_num' in infos.keys():
                    coop_num.append(infos['coop&coop_num'])
                    logger.add_scalars(
                        'coop&coop_num_per_episode',
                        {'coop&coop_num_per_episode': coop_num[episode]},
                        episode)
                if 'gore0_num' in infos.keys():
                    gore1_num.append(infos['gore0_num'])
                    logger.add_scalars(
                        'gore1_num_per_episode',
                        {'gore1_num_per_episode': gore1_num[episode]}, episode)
                if 'gore1_num' in infos.keys():
                    gore2_num.append(infos['gore1_num'])
                    logger.add_scalars(
                        'gore2_num_per_episode',
                        {'gore2_num_per_episode': gore2_num[episode]}, episode)
                if 'gore2_num' in infos.keys():
                    gore3_num.append(infos['gore2_num'])
                    logger.add_scalars(
                        'gore3_num_per_episode',
                        {'gore3_num_per_episode': gore3_num[episode]}, episode)
                if 'hare0_num' in infos.keys():
                    hare1_num.append(infos['hare0_num'])
                    logger.add_scalars(
                        'hare1_num_per_episode',
                        {'hare1_num_per_episode': hare1_num[episode]}, episode)
                if 'hare1_num' in infos.keys():
                    hare2_num.append(infos['hare1_num'])
                    logger.add_scalars(
                        'hare2_num_per_episode',
                        {'hare2_num_per_episode': hare2_num[episode]}, episode)
                if 'hare2_num' in infos.keys():
                    hare3_num.append(infos['hare2_num'])
                    logger.add_scalars(
                        'hare3_num_per_episode',
                        {'hare3_num_per_episode': hare3_num[episode]}, episode)
        logger.export_scalars_to_json(str(log_dir / 'summary.json'))
        logger.close()
コード例 #3
0
def main():
    args = get_config()

    # seed
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)

    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(args.n_training_threads)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)

    # path
    model_dir = Path(
        './results') / args.env_name / args.scenario_name / args.algorithm_name
    if not model_dir.exists():
        curr_run = 'run1'
    else:
        exst_run_nums = [
            int(str(folder.name).split('run')[1])
            for folder in model_dir.iterdir()
            if str(folder.name).startswith('run')
        ]
        if len(exst_run_nums) == 0:
            curr_run = 'run1'
        else:
            curr_run = 'run%i' % (max(exst_run_nums) + 1)

    run_dir = model_dir / curr_run
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir))

    # env
    envs = make_parallel_env(args)
    if args.eval:
        eval_env = make_eval_env(args)

    num_agents = args.num_agents
    all_action_space = []
    all_obs_space = []
    action_movement_dim = []
    if args.env_name == "BlueprintConstruction":
        order_obs = [
            'agent_qpos_qvel', 'box_obs', 'ramp_obs', 'construction_site_obs',
            'observation_self'
        ]
        mask_order_obs = [None, None, None, None, None]
    elif args.env_name == "BoxLocking":
        order_obs = [
            'agent_qpos_qvel', 'box_obs', 'ramp_obs', 'observation_self'
        ]
        mask_order_obs = ['mask_aa_obs', 'mask_ab_obs', 'mask_ar_obs', None]
    else:
        print("Can not support the " + args.env_name + "environment.")
        raise NotImplementedError

    for agent_id in range(num_agents):
        # deal with dict action space
        action_movement = envs.action_space['action_movement'][agent_id].nvec
        action_movement_dim.append(len(action_movement))
        action_glueall = envs.action_space['action_glueall'][agent_id].n
        action_vec = np.append(action_movement, action_glueall)
        if 'action_pull' in envs.action_space.spaces.keys():
            action_pull = envs.action_space['action_pull'][agent_id].n
            action_vec = np.append(action_vec, action_pull)
        action_space = MultiDiscrete([[0, vec - 1] for vec in action_vec])
        all_action_space.append(action_space)
        # deal with dict obs space
        obs_space = []
        obs_dim = 0
        for key in order_obs:
            if key in envs.observation_space.spaces.keys():
                space = list(envs.observation_space[key].shape)
                if len(space) < 2:
                    space.insert(0, 1)
                obs_space.append(space)
                obs_dim += reduce(lambda x, y: x * y, space)
        obs_space.insert(0, obs_dim)
        all_obs_space.append(obs_space)

    if args.share_policy:
        actor_critic = Policy(all_obs_space[0],
                              all_action_space[0],
                              num_agents=num_agents,
                              gain=args.gain,
                              base_kwargs={
                                  'naive_recurrent':
                                  args.naive_recurrent_policy,
                                  'recurrent': args.recurrent_policy,
                                  'hidden_size': args.hidden_size,
                                  'recurrent_N': args.recurrent_N,
                                  'attn': args.attn,
                                  'attn_only_critic': args.attn_only_critic,
                                  'attn_size': args.attn_size,
                                  'attn_N': args.attn_N,
                                  'attn_heads': args.attn_heads,
                                  'dropout': args.dropout,
                                  'use_average_pool': args.use_average_pool,
                                  'use_common_layer': args.use_common_layer,
                                  'use_feature_normlization':
                                  args.use_feature_normlization,
                                  'use_feature_popart':
                                  args.use_feature_popart,
                                  'use_orthogonal': args.use_orthogonal,
                                  'layer_N': args.layer_N,
                                  'use_ReLU': args.use_ReLU,
                                  'use_same_dim': True
                              },
                              device=device)
        actor_critic.to(device)
        # algorithm
        agents = PPO(actor_critic,
                     args.clip_param,
                     args.ppo_epoch,
                     args.num_mini_batch,
                     args.data_chunk_length,
                     args.value_loss_coef,
                     args.entropy_coef,
                     logger,
                     lr=args.lr,
                     eps=args.eps,
                     weight_decay=args.weight_decay,
                     max_grad_norm=args.max_grad_norm,
                     use_max_grad_norm=args.use_max_grad_norm,
                     use_clipped_value_loss=args.use_clipped_value_loss,
                     use_common_layer=args.use_common_layer,
                     use_huber_loss=args.use_huber_loss,
                     huber_delta=args.huber_delta,
                     use_popart=args.use_popart,
                     use_value_high_masks=args.use_value_high_masks,
                     device=device)

        #replay buffer
        rollouts = RolloutStorage(num_agents,
                                  args.episode_length,
                                  args.n_rollout_threads,
                                  all_obs_space[0],
                                  all_action_space[0],
                                  args.hidden_size,
                                  use_same_dim=True)
    else:
        actor_critic = []
        agents = []
        for agent_id in range(num_agents):
            ac = Policy(all_obs_space[0],
                        all_action_space[0],
                        num_agents=num_agents,
                        gain=args.gain,
                        base_kwargs={
                            'naive_recurrent': args.naive_recurrent_policy,
                            'recurrent': args.recurrent_policy,
                            'hidden_size': args.hidden_size,
                            'recurrent_N': args.recurrent_N,
                            'attn': args.attn,
                            'attn_only_critic': args.attn_only_critic,
                            'attn_size': args.attn_size,
                            'attn_N': args.attn_N,
                            'attn_heads': args.attn_heads,
                            'dropout': args.dropout,
                            'use_average_pool': args.use_average_pool,
                            'use_common_layer': args.use_common_layer,
                            'use_feature_normlization':
                            args.use_feature_normlization,
                            'use_feature_popart': args.use_feature_popart,
                            'use_orthogonal': args.use_orthogonal,
                            'layer_N': args.layer_N,
                            'use_ReLU': args.use_ReLU,
                            'use_same_dim': True
                        },
                        device=device)
            ac.to(device)
            # algorithm
            agent = PPO(ac,
                        args.clip_param,
                        args.ppo_epoch,
                        args.num_mini_batch,
                        args.data_chunk_length,
                        args.value_loss_coef,
                        args.entropy_coef,
                        logger,
                        lr=args.lr,
                        eps=args.eps,
                        weight_decay=args.weight_decay,
                        max_grad_norm=args.max_grad_norm,
                        use_max_grad_norm=args.use_max_grad_norm,
                        use_clipped_value_loss=args.use_clipped_value_loss,
                        use_common_layer=args.use_common_layer,
                        use_huber_loss=args.use_huber_loss,
                        huber_delta=args.huber_delta,
                        use_popart=args.use_popart,
                        use_value_high_masks=args.use_value_high_masks,
                        device=device)

            actor_critic.append(ac)
            agents.append(agent)

        #replay buffer
        rollouts = RolloutStorage(num_agents,
                                  args.episode_length,
                                  args.n_rollout_threads,
                                  all_obs_space[0],
                                  all_action_space[0],
                                  args.hidden_size,
                                  use_same_dim=True)

    # reset env
    dict_obs = envs.reset()
    obs = []
    share_obs = []
    for d_o in dict_obs:
        for i, key in enumerate(order_obs):
            if key in envs.observation_space.spaces.keys():
                if mask_order_obs[i] == None:
                    temp_share_obs = d_o[key].reshape(num_agents, -1).copy()
                    temp_obs = temp_share_obs.copy()
                else:
                    temp_share_obs = d_o[key].reshape(num_agents, -1).copy()
                    temp_mask = d_o[mask_order_obs[i]].copy()
                    temp_obs = d_o[key].copy()
                    mins_temp_mask = ~temp_mask
                    temp_obs[mins_temp_mask] = np.zeros(
                        (mins_temp_mask.sum(), temp_obs.shape[2]))
                    temp_obs = temp_obs.reshape(num_agents, -1)
                if i == 0:
                    reshape_obs = temp_obs.copy()
                    reshape_share_obs = temp_share_obs.copy()
                else:
                    reshape_obs = np.concatenate((reshape_obs, temp_obs),
                                                 axis=1)
                    reshape_share_obs = np.concatenate(
                        (reshape_share_obs, temp_share_obs), axis=1)
        obs.append(reshape_obs)
        share_obs.append(reshape_share_obs)
    obs = np.array(obs)
    share_obs = np.array(share_obs)

    # replay buffer
    rollouts.share_obs[0] = share_obs.copy()
    rollouts.obs[0] = obs.copy()
    rollouts.recurrent_hidden_states = np.zeros(
        rollouts.recurrent_hidden_states.shape).astype(np.float32)
    rollouts.recurrent_hidden_states_critic = np.zeros(
        rollouts.recurrent_hidden_states_critic.shape).astype(np.float32)

    # run
    start = time.time()
    episodes = int(
        args.num_env_steps) // args.episode_length // args.n_rollout_threads
    timesteps = 0

    for episode in range(episodes):
        if args.use_linear_lr_decay:  # decrease learning rate linearly
            if args.share_policy:
                update_linear_schedule(agents.optimizer, episode, episodes,
                                       args.lr)
            else:
                for agent_id in range(num_agents):
                    update_linear_schedule(agents[agent_id].optimizer, episode,
                                           episodes, args.lr)
        # info list
        discard_episode = 0
        success = 0
        trials = 0

        for step in range(args.episode_length):
            # Sample actions
            values = []
            actions = []
            action_log_probs = []
            recurrent_hidden_statess = []
            recurrent_hidden_statess_critic = []
            with torch.no_grad():
                for agent_id in range(num_agents):
                    if args.share_policy:
                        actor_critic.eval()
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(
                            agent_id,
                            torch.tensor(rollouts.share_obs[step, :,
                                                            agent_id]),
                            torch.tensor(rollouts.obs[step, :, agent_id]),
                            torch.tensor(
                                rollouts.recurrent_hidden_states[step, :,
                                                                 agent_id]),
                            torch.tensor(
                                rollouts.recurrent_hidden_states_critic[
                                    step, :, agent_id]),
                            torch.tensor(rollouts.masks[step, :, agent_id]))
                    else:
                        actor_critic[agent_id].eval()
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic[
                            agent_id].act(
                                agent_id,
                                torch.tensor(rollouts.share_obs[step, :,
                                                                agent_id]),
                                torch.tensor(rollouts.obs[step, :, agent_id]),
                                torch.tensor(rollouts.recurrent_hidden_states[
                                    step, :, agent_id]),
                                torch.tensor(
                                    rollouts.recurrent_hidden_states_critic[
                                        step, :, agent_id]),
                                torch.tensor(rollouts.masks[step, :,
                                                            agent_id]))

                    values.append(value.detach().cpu().numpy())
                    actions.append(action.detach().cpu().numpy())
                    action_log_probs.append(
                        action_log_prob.detach().cpu().numpy())
                    recurrent_hidden_statess.append(
                        recurrent_hidden_states.detach().cpu().numpy())
                    recurrent_hidden_statess_critic.append(
                        recurrent_hidden_states_critic.detach().cpu().numpy())

            # rearrange action
            actions_env = []
            for n_rollout_thread in range(args.n_rollout_threads):
                action_movement = []
                action_pull = []
                action_glueall = []
                for agent_id in range(num_agents):
                    action_movement.append(actions[agent_id][n_rollout_thread]
                                           [:action_movement_dim[agent_id]])
                    action_glueall.append(
                        int(actions[agent_id][n_rollout_thread][
                            action_movement_dim[agent_id]]))
                    if 'action_pull' in envs.action_space.spaces.keys():
                        action_pull.append(
                            int(actions[agent_id][n_rollout_thread][-1]))
                action_movement = np.stack(action_movement, axis=0)
                action_glueall = np.stack(action_glueall, axis=0)
                if 'action_pull' in envs.action_space.spaces.keys():
                    action_pull = np.stack(action_pull, axis=0)
                one_env_action = {
                    'action_movement': action_movement,
                    'action_pull': action_pull,
                    'action_glueall': action_glueall
                }
                actions_env.append(one_env_action)

            # Obser reward and next obs
            dict_obs, rewards, dones, infos = envs.step(actions_env)
            if len(rewards.shape) < 3:
                rewards = rewards[:, :, np.newaxis]

            # If done then clean the history of observations.
            # insert data in buffer
            masks = []
            for i, done in enumerate(dones):
                if done:
                    if "discard_episode" in infos[i].keys():
                        if infos[i]['discard_episode']:
                            discard_episode += 1
                        else:
                            trials += 1
                    else:
                        trials += 1
                    if "success" in infos[i].keys():
                        if infos[i]['success']:
                            success += 1
                mask = []
                for agent_id in range(num_agents):
                    if done:
                        recurrent_hidden_statess[agent_id][i] = np.zeros(
                            args.hidden_size).astype(np.float32)
                        recurrent_hidden_statess_critic[agent_id][
                            i] = np.zeros(args.hidden_size).astype(np.float32)
                        mask.append([0.0])
                    else:
                        mask.append([1.0])
                masks.append(mask)

            obs = []
            share_obs = []
            for d_o in dict_obs:
                for i, key in enumerate(order_obs):
                    if key in envs.observation_space.spaces.keys():
                        if mask_order_obs[i] == None:
                            temp_share_obs = d_o[key].reshape(num_agents,
                                                              -1).copy()
                            temp_obs = temp_share_obs.copy()
                        else:
                            temp_share_obs = d_o[key].reshape(num_agents,
                                                              -1).copy()
                            temp_mask = d_o[mask_order_obs[i]].copy()
                            temp_obs = d_o[key].copy()
                            mins_temp_mask = ~temp_mask
                            temp_obs[mins_temp_mask] = np.zeros(
                                (mins_temp_mask.sum(), temp_obs.shape[2]))
                            temp_obs = temp_obs.reshape(num_agents, -1)
                        if i == 0:
                            reshape_obs = temp_obs.copy()
                            reshape_share_obs = temp_share_obs.copy()
                        else:
                            reshape_obs = np.concatenate(
                                (reshape_obs, temp_obs), axis=1)
                            reshape_share_obs = np.concatenate(
                                (reshape_share_obs, temp_share_obs), axis=1)
                obs.append(reshape_obs)
                share_obs.append(reshape_share_obs)
            obs = np.array(obs)
            share_obs = np.array(share_obs)

            rollouts.insert(
                share_obs, obs,
                np.array(recurrent_hidden_statess).transpose(1, 0, 2),
                np.array(recurrent_hidden_statess_critic).transpose(1, 0, 2),
                np.array(actions).transpose(1, 0, 2),
                np.array(action_log_probs).transpose(1, 0, 2),
                np.array(values).transpose(1, 0, 2), rewards, masks)

        with torch.no_grad():
            for agent_id in range(num_agents):
                if args.share_policy:
                    actor_critic.eval()
                    next_value, _, _ = actor_critic.get_value(
                        agent_id,
                        torch.tensor(rollouts.share_obs[-1, :, agent_id]),
                        torch.tensor(rollouts.obs[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states_critic[-1, :,
                                                                    agent_id]),
                        torch.tensor(rollouts.masks[-1, :, agent_id]))
                    next_value = next_value.detach().cpu().numpy()
                    rollouts.compute_returns(agent_id, next_value,
                                             args.use_gae, args.gamma,
                                             args.gae_lambda,
                                             args.use_proper_time_limits,
                                             args.use_popart,
                                             agents.value_normalizer)
                else:
                    actor_critic[agent_id].eval()
                    next_value, _, _ = actor_critic[agent_id].get_value(
                        agent_id,
                        torch.tensor(rollouts.share_obs[-1, :, agent_id]),
                        torch.tensor(rollouts.obs[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states_critic[-1, :,
                                                                    agent_id]),
                        torch.tensor(rollouts.masks[-1, :, agent_id]))
                    next_value = next_value.detach().cpu().numpy()
                    rollouts.compute_returns(agent_id, next_value,
                                             args.use_gae, args.gamma,
                                             args.gae_lambda,
                                             args.use_proper_time_limits,
                                             args.use_popart,
                                             agents[agent_id].value_normalizer)

        # update the network
        if args.share_policy:
            actor_critic.train()
            value_loss, action_loss, dist_entropy = agents.update_share(
                num_agents, rollouts)

            logger.add_scalars('reward', {'reward': np.mean(rollouts.rewards)},
                               (episode + 1) * args.episode_length *
                               args.n_rollout_threads)
        else:
            value_losses = []
            action_losses = []
            dist_entropies = []

            for agent_id in range(num_agents):
                actor_critic[agent_id].train()
                value_loss, action_loss, dist_entropy = agents[
                    agent_id].update(agent_id, rollouts)
                value_losses.append(value_loss)
                action_losses.append(action_loss)
                dist_entropies.append(dist_entropy)

                logger.add_scalars(
                    'agent%i/reward' % agent_id,
                    {'reward': np.mean(rollouts.rewards[:, :, agent_id])},
                    (episode + 1) * args.episode_length *
                    args.n_rollout_threads)

        # clean the buffer and reset
        rollouts.after_update()

        total_num_steps = (episode +
                           1) * args.episode_length * args.n_rollout_threads

        if (episode % args.save_interval == 0 or episode == episodes -
                1):  # save for every interval-th episode or for the last epoch
            if args.share_policy:
                torch.save({'model': actor_critic},
                           str(save_dir) + "/agent_model.pt")
            else:
                for agent_id in range(num_agents):
                    torch.save({'model': actor_critic[agent_id]},
                               str(save_dir) + "/agent%i_model" % agent_id +
                               ".pt")

        # log information
        if episode % args.log_interval == 0:
            end = time.time()
            print(
                "\n Scenario {} Algo {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(args.scenario_name, args.algorithm_name, episode,
                        episodes, total_num_steps, args.num_env_steps,
                        int(total_num_steps / (end - start))))
            if args.share_policy:
                print("value loss of agent: " + str(value_loss))
            else:
                for agent_id in range(num_agents):
                    print("value loss of agent%i: " % agent_id +
                          str(value_losses[agent_id]))

            logger.add_scalars('discard_episode',
                               {'discard_episode': discard_episode},
                               total_num_steps)
            if trials > 0:
                logger.add_scalars('success_rate',
                                   {'success_rate': success / trials},
                                   total_num_steps)
            else:
                logger.add_scalars('success_rate', {'success_rate': 0.0},
                                   total_num_steps)
        # eval
        if episode % args.eval_interval == 0 and args.eval:
            eval_episode = 0
            eval_success = 0
            eval_dict_obs = eval_env.reset()

            eval_obs = []
            eval_share_obs = []
            for eval_d_o in eval_dict_obs:
                for i, key in enumerate(order_obs):
                    if key in eval_env.observation_space.spaces.keys():
                        if mask_order_obs[i] == None:
                            temp_share_obs = eval_d_o[key].reshape(
                                num_agents, -1).copy()
                            temp_obs = temp_share_obs.copy()
                        else:
                            temp_share_obs = eval_d_o[key].reshape(
                                num_agents, -1).copy()
                            temp_mask = eval_d_o[mask_order_obs[i]].copy()
                            temp_obs = eval_d_o[key].copy()
                            mins_temp_mask = ~temp_mask
                            temp_obs[mins_temp_mask] = np.zeros(
                                (mins_temp_mask.sum(), temp_obs.shape[2]))
                            temp_obs = temp_obs.reshape(num_agents, -1)
                        if i == 0:
                            reshape_obs = temp_obs.copy()
                            reshape_share_obs = temp_share_obs.copy()
                        else:
                            reshape_obs = np.concatenate(
                                (reshape_obs, temp_obs), axis=1)
                            reshape_share_obs = np.concatenate(
                                (reshape_share_obs, temp_share_obs), axis=1)
                eval_obs.append(reshape_obs)
                eval_share_obs.append(reshape_share_obs)
            eval_obs = np.array(eval_obs)
            eval_share_obs = np.array(eval_share_obs)

            eval_recurrent_hidden_states = np.zeros(
                (1, num_agents, args.hidden_size)).astype(np.float32)
            eval_recurrent_hidden_states_critic = np.zeros(
                (1, num_agents, args.hidden_size)).astype(np.float32)
            eval_masks = np.ones((1, num_agents, 1)).astype(np.float32)

            while True:
                eval_actions = []
                actor_critic.eval()
                for agent_id in range(num_agents):
                    _, action, _, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(
                        agent_id,
                        torch.FloatTensor(eval_share_obs[:, agent_id]),
                        torch.FloatTensor(eval_obs[:, agent_id]),
                        torch.FloatTensor(
                            eval_recurrent_hidden_states[:, agent_id]),
                        torch.FloatTensor(
                            eval_recurrent_hidden_states_critic[:, agent_id]),
                        torch.FloatTensor(eval_masks[:, agent_id]),
                        None,
                        deterministic=True)

                    eval_actions.append(action.detach().cpu().numpy())
                    eval_recurrent_hidden_states[:,
                                                 agent_id] = recurrent_hidden_states.detach(
                                                 ).cpu().numpy()
                    eval_recurrent_hidden_states_critic[:,
                                                        agent_id] = recurrent_hidden_states_critic.detach(
                                                        ).cpu().numpy()

                # rearrange action
                eval_actions_env = []
                for n_rollout_thread in range(1):
                    action_movement = []
                    action_pull = []
                    action_glueall = []
                    for agent_id in range(num_agents):
                        action_movement.append(
                            eval_actions[agent_id][n_rollout_thread]
                            [:action_movement_dim[agent_id]])
                        action_glueall.append(
                            int(eval_actions[agent_id][n_rollout_thread][
                                action_movement_dim[agent_id]]))
                        if 'action_pull' in envs.action_space.spaces.keys():
                            action_pull.append(
                                int(eval_actions[agent_id][n_rollout_thread]
                                    [-1]))
                    action_movement = np.stack(action_movement, axis=0)
                    action_glueall = np.stack(action_glueall, axis=0)
                    if 'action_pull' in envs.action_space.spaces.keys():
                        action_pull = np.stack(action_pull, axis=0)
                    one_env_action = {
                        'action_movement': action_movement,
                        'action_pull': action_pull,
                        'action_glueall': action_glueall
                    }
                    eval_actions_env.append(one_env_action)

                # Obser reward and next obs
                eval_dict_obs, eval_rewards, eval_dones, eval_infos = eval_env.step(
                    eval_actions_env)

                eval_obs = []
                eval_share_obs = []
                for eval_d_o in eval_dict_obs:
                    for i, key in enumerate(order_obs):
                        if key in eval_env.observation_space.spaces.keys():
                            if mask_order_obs[i] == None:
                                temp_share_obs = eval_d_o[key].reshape(
                                    num_agents, -1).copy()
                                temp_obs = temp_share_obs.copy()
                            else:
                                temp_share_obs = eval_d_o[key].reshape(
                                    num_agents, -1).copy()
                                temp_mask = eval_d_o[mask_order_obs[i]].copy()
                                temp_obs = eval_d_o[key].copy()
                                mins_temp_mask = ~temp_mask
                                temp_obs[mins_temp_mask] = np.zeros(
                                    (mins_temp_mask.sum(), temp_obs.shape[2]))
                                temp_obs = temp_obs.reshape(num_agents, -1)
                            if i == 0:
                                reshape_obs = temp_obs.copy()
                                reshape_share_obs = temp_share_obs.copy()
                            else:
                                reshape_obs = np.concatenate(
                                    (reshape_obs, temp_obs), axis=1)
                                reshape_share_obs = np.concatenate(
                                    (reshape_share_obs, temp_share_obs),
                                    axis=1)
                    eval_obs.append(reshape_obs)
                    eval_share_obs.append(reshape_share_obs)
                eval_obs = np.array(eval_obs)
                eval_share_obs = np.array(eval_share_obs)

                eval_recurrent_hidden_states = np.zeros(
                    (1, num_agents, args.hidden_size)).astype(np.float32)
                eval_recurrent_hidden_states_critic = np.zeros(
                    (1, num_agents, args.hidden_size)).astype(np.float32)
                eval_masks = np.ones((1, num_agents, 1)).astype(np.float32)

                if eval_dones[0]:
                    eval_episode += 1
                    if "success" in eval_infos[0].keys():
                        if eval_infos[0]['success']:
                            eval_success += 1
                    for agent_id in range(num_agents):
                        eval_recurrent_hidden_states[0][agent_id] = np.zeros(
                            args.hidden_size).astype(np.float32)
                        eval_recurrent_hidden_states_critic[0][
                            agent_id] = np.zeros(args.hidden_size).astype(
                                np.float32)
                        eval_masks[0][agent_id] = 0.0
                else:
                    for agent_id in range(num_agents):
                        eval_masks[0][agent_id] = 1.0

                if eval_episode >= args.eval_episodes:
                    logger.add_scalars('eval_success_rate', {
                        'eval_success_rate':
                        eval_success / args.eval_episodes
                    }, total_num_steps)
                    print("eval_success_rate is " +
                          str(eval_success / args.eval_episodes))
                    break

    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()
    envs.close()
    if args.eval:
        eval_env.close()
コード例 #4
0
ファイル: train_mpe.py プロジェクト: zhangyx96/ma_goal_gan
def main():
    args = get_config()

    assert (
        args.share_policy == True
        and args.scenario_name == 'simple_speaker_listener'
    ) == False, (
        "The simple_speaker_listener scenario can not use shared policy. Please check the config.py."
    )

    # seed
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(1)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)

    # path
    model_dir = Path(
        './results') / args.env_name / args.scenario_name / args.algorithm_name
    if not model_dir.exists():
        curr_run = 'run1'
    else:
        exst_run_nums = [
            int(str(folder.name).split('run')[1])
            for folder in model_dir.iterdir()
            if str(folder.name).startswith('run')
        ]
        if len(exst_run_nums) == 0:
            curr_run = 'run1'
        else:
            curr_run = 'run%i' % (max(exst_run_nums) + 1)

    run_dir = model_dir / curr_run
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir))

    # env
    envs = make_parallel_env(args)
    num_agents = args.num_agents
    #Policy network
    if args.share_policy:
        share_base = ATTBase_add(envs.observation_space[0].shape[0],
                                 num_agents)
        actor_critic = Policy(envs.observation_space[0],
                              envs.action_space[0],
                              num_agents=num_agents,
                              base=share_base,
                              base_kwargs={
                                  'naive_recurrent':
                                  args.naive_recurrent_policy,
                                  'recurrent': args.recurrent_policy,
                                  'hidden_size': args.hidden_size,
                                  'attn': args.attn,
                                  'attn_size': args.attn_size,
                                  'attn_N': args.attn_N,
                                  'attn_heads': args.attn_heads,
                                  'dropout': args.dropout,
                                  'use_average_pool': args.use_average_pool,
                                  'use_common_layer': args.use_common_layer,
                                  'use_feature_normlization':
                                  args.use_feature_normlization,
                                  'use_feature_popart':
                                  args.use_feature_popart,
                                  'use_orthogonal': args.use_orthogonal,
                                  'layer_N': args.layer_N,
                                  'use_ReLU': args.use_ReLU
                              },
                              device=device)
        actor_critic.to(device)
        # algorithm
        agents = PPO(actor_critic,
                     args.clip_param,
                     args.ppo_epoch,
                     args.num_mini_batch,
                     args.data_chunk_length,
                     args.value_loss_coef,
                     args.entropy_coef,
                     logger,
                     lr=args.lr,
                     eps=args.eps,
                     weight_decay=args.weight_decay,
                     max_grad_norm=args.max_grad_norm,
                     use_max_grad_norm=args.use_max_grad_norm,
                     use_clipped_value_loss=args.use_clipped_value_loss,
                     use_common_layer=args.use_common_layer,
                     use_huber_loss=args.use_huber_loss,
                     huber_delta=args.huber_delta,
                     use_popart=args.use_popart,
                     device=device)

        #replay buffer
        rollouts = RolloutStorage_share(num_agents, args.episode_length,
                                        args.n_rollout_threads,
                                        envs.observation_space[0],
                                        envs.action_space[0], args.hidden_size)
    else:
        actor_critic = []
        agents = []
        rollouts = []
        for agent_id in range(num_agents):
            ac = Policy(
                envs.observation_space,
                envs.action_space[agent_id],
                num_agents=agent_id,  # here is special
                base_kwargs={
                    'naive_recurrent': args.naive_recurrent_policy,
                    'recurrent': args.recurrent_policy,
                    'hidden_size': args.hidden_size,
                    'attn': args.attn,
                    'attn_size': args.attn_size,
                    'attn_N': args.attn_N,
                    'attn_heads': args.attn_heads,
                    'dropout': args.dropout,
                    'use_average_pool': args.use_average_pool,
                    'use_common_layer': args.use_common_layer,
                    'use_feature_normlization': args.use_feature_normlization,
                    'use_feature_popart': args.use_feature_popart,
                    'use_orthogonal': args.use_orthogonal,
                    'layer_N': args.layer_N,
                    'use_ReLU': args.use_ReLU
                },
                device=device)
            ac.to(device)
            # algorithm
            agent = PPO(ac,
                        args.clip_param,
                        args.ppo_epoch,
                        args.num_mini_batch,
                        args.data_chunk_length,
                        args.value_loss_coef,
                        args.entropy_coef,
                        logger,
                        lr=args.lr,
                        eps=args.eps,
                        weight_decay=args.weight_decay,
                        max_grad_norm=args.max_grad_norm,
                        use_max_grad_norm=args.use_max_grad_norm,
                        use_clipped_value_loss=args.use_clipped_value_loss,
                        use_common_layer=args.use_common_layer,
                        use_huber_loss=args.use_huber_loss,
                        huber_delta=args.huber_delta,
                        use_popart=args.use_popart,
                        device=device)

            actor_critic.append(ac)
            agents.append(agent)

            #replay buffer
            ro = SingleRolloutStorage(agent_id, args.episode_length,
                                      args.n_rollout_threads,
                                      envs.observation_space,
                                      envs.action_space, args.hidden_size)
            rollouts.append(ro)

    # reset env
    obs, _ = envs.reset(num_agents)

    # replay buffer
    if args.share_policy:
        share_obs = obs.reshape(args.n_rollout_threads, -1)
        # share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)
        rollouts.share_obs[0] = share_obs.copy()
        rollouts.obs[0] = obs.copy()
        rollouts.recurrent_hidden_states = np.zeros(
            rollouts.recurrent_hidden_states.shape).astype(np.float32)
        rollouts.recurrent_hidden_states_critic = np.zeros(
            rollouts.recurrent_hidden_states_critic.shape).astype(np.float32)
    else:

        share_obs = []
        for o in obs:
            share_obs.append(list(itertools.chain(*o)))
        share_obs = np.array(share_obs)
        for agent_id in range(num_agents):
            rollouts[agent_id].share_obs[0] = share_obs.copy()
            rollouts[agent_id].obs[0] = np.array(list(obs[:, agent_id])).copy()
            rollouts[agent_id].recurrent_hidden_states = np.zeros(
                rollouts[agent_id].recurrent_hidden_states.shape).astype(
                    np.float32)
            rollouts[agent_id].recurrent_hidden_states_critic = np.zeros(
                rollouts[agent_id].recurrent_hidden_states_critic.shape
            ).astype(np.float32)

    # run
    start = time.time()
    episodes = int(
        args.num_env_steps) // args.episode_length // args.n_rollout_threads
    timesteps = 0

    for episode in range(episodes):
        if args.use_linear_lr_decay:  # decrease learning rate linearly
            if args.share_policy:
                update_linear_schedule(agents.optimizer, episode, episodes,
                                       args.lr)
            else:
                for agent_id in range(num_agents):
                    update_linear_schedule(agents[agent_id].optimizer, episode,
                                           episodes, args.lr)

        step_cover_rate = np.zeros(shape=(args.n_rollout_threads,
                                          args.episode_length))
        for step in range(args.episode_length):
            # Sample actions
            values = []
            actions = []
            action_log_probs = []
            recurrent_hidden_statess = []
            recurrent_hidden_statess_critic = []

            with torch.no_grad():
                for agent_id in range(num_agents):
                    if args.share_policy:
                        actor_critic.eval()
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(
                            agent_id,
                            # torch.FloatTensor(rollouts.share_obs[step,:,agent_id]),
                            torch.FloatTensor(rollouts.share_obs[step]),
                            torch.FloatTensor(rollouts.obs[step, :, agent_id]),
                            torch.FloatTensor(
                                rollouts.recurrent_hidden_states[step, :,
                                                                 agent_id]),
                            torch.FloatTensor(
                                rollouts.recurrent_hidden_states_critic[
                                    step, :, agent_id]),
                            torch.FloatTensor(rollouts.masks[step, :,
                                                             agent_id]))
                    else:
                        actor_critic[agent_id].eval()
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic[
                            agent_id].act(
                                agent_id,
                                torch.FloatTensor(
                                    rollouts[agent_id].share_obs[step, :]),
                                torch.FloatTensor(
                                    rollouts[agent_id].obs[step, :]),
                                torch.FloatTensor(
                                    rollouts[agent_id].recurrent_hidden_states[
                                        step, :]),
                                torch.FloatTensor(
                                    rollouts[agent_id].
                                    recurrent_hidden_states_critic[step, :]),
                                torch.FloatTensor(
                                    rollouts[agent_id].masks[step, :]))

                    values.append(value.detach().cpu().numpy())
                    actions.append(action.detach().cpu().numpy())
                    action_log_probs.append(
                        action_log_prob.detach().cpu().numpy())
                    recurrent_hidden_statess.append(
                        recurrent_hidden_states.detach().cpu().numpy())
                    recurrent_hidden_statess_critic.append(
                        recurrent_hidden_states_critic.detach().cpu().numpy())

            # rearrange action
            actions_env = []
            for i in range(args.n_rollout_threads):
                one_hot_action_env = []
                for agent_id in range(num_agents):
                    if envs.action_space[
                            agent_id].__class__.__name__ == 'MultiDiscrete':
                        uc_action = []
                        for j in range(envs.action_space[agent_id].shape):
                            uc_one_hot_action = np.zeros(
                                envs.action_space[agent_id].high[j] + 1)
                            uc_one_hot_action[actions[agent_id][i][j]] = 1
                            uc_action.append(uc_one_hot_action)
                        uc_action = np.concatenate(uc_action)
                        one_hot_action_env.append(uc_action)

                    elif envs.action_space[
                            agent_id].__class__.__name__ == 'Discrete':
                        one_hot_action = np.zeros(
                            envs.action_space[agent_id].n)
                        one_hot_action[actions[agent_id][i]] = 1
                        one_hot_action_env.append(one_hot_action)
                    else:
                        raise NotImplementedError
                actions_env.append(one_hot_action_env)

            # Obser reward and next obs
            obs, rewards, dones, infos, _ = envs.step(actions_env,
                                                      args.n_rollout_threads,
                                                      num_agents)
            step_cover_rate[:, step] = np.array(infos)[:, 0]

            # If done then clean the history of observations.
            # insert data in buffer
            masks = []
            for i, done in enumerate(dones):
                mask = []
                for agent_id in range(num_agents):
                    if done[agent_id]:
                        recurrent_hidden_statess[agent_id][i] = np.zeros(
                            args.hidden_size).astype(np.float32)
                        recurrent_hidden_statess_critic[agent_id][
                            i] = np.zeros(args.hidden_size).astype(np.float32)
                        mask.append([0.0])
                    else:
                        mask.append([1.0])
                masks.append(mask)

            if args.share_policy:
                share_obs = obs.reshape(args.n_rollout_threads, -1)
                # share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)

                rollouts.insert(
                    share_obs, obs,
                    np.array(recurrent_hidden_statess).transpose(1, 0, 2),
                    np.array(recurrent_hidden_statess_critic).transpose(
                        1, 0, 2),
                    np.array(actions).transpose(1, 0, 2),
                    np.array(action_log_probs).transpose(1, 0, 2),
                    np.array(values).transpose(1, 0, 2), rewards, masks)
            else:
                share_obs = []
                for o in obs:
                    share_obs.append(list(itertools.chain(*o)))
                share_obs = np.array(share_obs)
                for agent_id in range(num_agents):
                    rollouts[agent_id].insert(
                        share_obs, np.array(list(obs[:, agent_id])),
                        np.array(recurrent_hidden_statess[agent_id]),
                        np.array(recurrent_hidden_statess_critic[agent_id]),
                        np.array(actions[agent_id]),
                        np.array(action_log_probs[agent_id]),
                        np.array(values[agent_id]), rewards[:, agent_id],
                        np.array(masks)[:, agent_id])
        # import pdb;pdb.set_trace()
        logger.add_scalars(
            'agent/cover_rate_1step',
            {'cover_rate_1step': np.mean(step_cover_rate[:, -1])},
            (episode + 1) * args.episode_length * args.n_rollout_threads)
        logger.add_scalars('agent/cover_rate_5step', {
            'cover_rate_5step':
            np.mean(np.mean(step_cover_rate[:, -5:], axis=1))
        }, (episode + 1) * args.episode_length * args.n_rollout_threads)

        with torch.no_grad():
            for agent_id in range(num_agents):
                if args.share_policy:
                    actor_critic.eval()
                    next_value, _, _ = actor_critic.get_value(
                        agent_id,
                        #    torch.FloatTensor(rollouts.share_obs[-1,:,agent_id]),
                        torch.FloatTensor(rollouts.share_obs[-1]),
                        torch.FloatTensor(rollouts.obs[-1, :, agent_id]),
                        torch.FloatTensor(
                            rollouts.recurrent_hidden_states[-1, :, agent_id]),
                        torch.FloatTensor(
                            rollouts.recurrent_hidden_states_critic[-1, :,
                                                                    agent_id]),
                        torch.FloatTensor(rollouts.masks[-1, :, agent_id]))
                    next_value = next_value.detach().cpu().numpy()
                    rollouts.compute_returns(agent_id, next_value,
                                             args.use_gae, args.gamma,
                                             args.gae_lambda,
                                             args.use_proper_time_limits,
                                             args.use_popart,
                                             agents.value_normalizer)
                else:
                    actor_critic[agent_id].eval()
                    next_value, _, _ = actor_critic[agent_id].get_value(
                        agent_id,
                        torch.FloatTensor(rollouts[agent_id].share_obs[-1, :]),
                        torch.FloatTensor(rollouts[agent_id].obs[-1, :]),
                        torch.FloatTensor(
                            rollouts[agent_id].recurrent_hidden_states[-1, :]),
                        torch.FloatTensor(
                            rollouts[agent_id].recurrent_hidden_states_critic[
                                -1, :]),
                        torch.FloatTensor(rollouts[agent_id].masks[-1, :]))
                    next_value = next_value.detach().cpu().numpy()
                    rollouts[agent_id].compute_returns(
                        next_value, args.use_gae, args.gamma, args.gae_lambda,
                        args.use_proper_time_limits, args.use_popart,
                        agents[agent_id].value_normalizer)

        # update the network
        if args.share_policy:
            actor_critic.train()
            value_loss, action_loss, dist_entropy = agents.update_share(
                num_agents, rollouts)

            rew = []
            for i in range(rollouts.rewards.shape[1]):
                rew.append(np.sum(rollouts.rewards[:, i]))
            logger.add_scalars('average_episode_reward',
                               {'average_episode_reward': np.mean(rew)},
                               (episode + 1) * args.episode_length *
                               args.n_rollout_threads)
            # clean the buffer and reset
            rollouts.after_update()
        else:
            value_losses = []
            action_losses = []
            dist_entropies = []

            for agent_id in range(num_agents):
                actor_critic[agent_id].train()
                value_loss, action_loss, dist_entropy = agents[
                    agent_id].update_single(agent_id, rollouts[agent_id])
                value_losses.append(value_loss)
                action_losses.append(action_loss)
                dist_entropies.append(dist_entropy)

                rew = []
                for i in range(rollouts[agent_id].rewards.shape[1]):
                    rew.append(np.sum(rollouts[agent_id].rewards[:, i]))
                logger.add_scalars('agent%i/average_episode_reward' % agent_id,
                                   {'average_episode_reward': np.mean(rew)},
                                   (episode + 1) * args.episode_length *
                                   args.n_rollout_threads)

                rollouts[agent_id].after_update()

        obs, _ = envs.reset(num_agents)
        # replay buffer
        if args.share_policy:
            share_obs = obs.reshape(args.n_rollout_threads, -1)
            # share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)
            rollouts.share_obs[0] = share_obs.copy()
            rollouts.obs[0] = obs.copy()
            rollouts.recurrent_hidden_states = np.zeros(
                rollouts.recurrent_hidden_states.shape).astype(np.float32)
            rollouts.recurrent_hidden_states_critic = np.zeros(
                rollouts.recurrent_hidden_states_critic.shape).astype(
                    np.float32)
        else:

            share_obs = []
            for o in obs:
                share_obs.append(list(itertools.chain(*o)))
            share_obs = np.array(share_obs)
            for agent_id in range(num_agents):
                rollouts[agent_id].share_obs[0] = share_obs.copy()
                rollouts[agent_id].obs[0] = np.array(list(
                    obs[:, agent_id])).copy()
                rollouts[agent_id].recurrent_hidden_states = np.zeros(
                    rollouts[agent_id].recurrent_hidden_states.shape).astype(
                        np.float32)
                rollouts[agent_id].recurrent_hidden_states_critic = np.zeros(
                    rollouts[agent_id].recurrent_hidden_states_critic.shape
                ).astype(np.float32)

        total_num_steps = (episode +
                           1) * args.episode_length * args.n_rollout_threads

        if (episode % args.save_interval == 0 or episode == episodes -
                1):  # save for every interval-th episode or for the last epoch
            if args.share_policy:
                torch.save({'model': actor_critic},
                           str(save_dir) + "/agent_model.pt")
            else:
                for agent_id in range(num_agents):
                    torch.save({'model': actor_critic[agent_id]},
                               str(save_dir) + "/agent%i_model" % agent_id +
                               ".pt")

        # log information
        if episode % args.log_interval == 0:
            end = time.time()
            print(
                "\n Scenario {} Algo {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(args.scenario_name, args.algorithm_name, episode,
                        episodes, total_num_steps, args.num_env_steps,
                        int(total_num_steps / (end - start))))
            if args.share_policy:
                print("value loss of agent: " + str(value_loss))
            else:
                for agent_id in range(num_agents):
                    print("value loss of agent%i: " % agent_id +
                          str(value_losses[agent_id]))

            # if args.env_name == "MPE":
            #     for agent_id in range(num_agents):
            #         show_rewards = []
            #         for info in infos:
            #             if 'individual_reward' in info[agent_id].keys():
            #                 show_rewards.append(info[agent_id]['individual_reward'])
            #         logger.add_scalars('agent%i/individual_reward' % agent_id, {'individual_reward': np.mean(show_rewards)}, total_num_steps)

    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()
    envs.close()
コード例 #5
0
class Trainer:
    def __init__(self, args):
        if args.seed is not None:
            torch.manual_seed(args.seed)
            np.random.seed(args.seed)
            random.seed(args.seed)

        self.env_name = args.environment
        self.env_setting = get_env_setting(self.env_name)
        self.solved_reward = self.env_setting["solved_reward"]
        self.update_timestep = self.env_setting["update_timestep"]
        self.device = torch.device(
            "cuda:0" if torch.cuda.is_available() else "cpu")
        self.env = gym.make(args.environment)
        self.alg = PPO(args, self.device)
        self.log_interval = 5  # print avg reward in the interval
        self.max_episodes = 100000
        self.render = False

    def train(self):
        # logging variables
        running_reward = 0
        avg_length = 0
        time_step = 0
        memory = Memory()
        # self.alg.load_dict("./", self.env_name, self.alg_name, self.net_name)

        # training loop
        time_step = 0
        for i_episode in range(1, self.max_episodes + 1):
            self.alg.reset_memory()
            obs = self.env.reset(render_mode="logic")
            t = 0
            while True:
                t += 1
                # Running policy_old:
                action = self.alg.take_action(obs, memory)
                self.env.render()
                obs, reward, done, _ = self.env.step(action,
                                                     observation_mode="logic")

                # Saving reward and is_terminal:
                memory.rewards.append(reward)
                memory.is_terminals.append(done)

                running_reward += reward
                if self.render:
                    self.env.render()
                if done:
                    break
            time_step += t

            # update if its time
            if time_step >= self.update_timestep and done == True:
                self.alg.update(memory)
                memory.clear_memory()
                time_step = 0

            avg_length += t

            # save every 500 episodes
            if i_episode % 500 == 0:
                directory = "./epoch_performance"
                if not os.path.exists(directory):
                    os.makedirs(directory)
                self.alg.save_dict(directory, f'{self.env_name}_{i_episode}')

            # logging
            if i_episode % self.log_interval == 0:
                avg_length = int(avg_length / self.log_interval)
                running_reward = int((running_reward / self.log_interval))
                print('Episode {} \t avg length: {} \t reward: {}'.format(
                    i_episode, avg_length, running_reward))

                # stop training if avg_reward > solved_reward or reaches the limit of training epoches
                if running_reward > (self.log_interval * self.solved_reward):
                    print("########## Solved! ##########")
                    directory = "./success/"
                    if not os.path.exists(directory):
                        os.makedirs(directory)
                    self.alg.save_dict(directory,
                                       f'{self.env_name}_{self.log_interval}')
                    break

                running_reward = 0
                avg_length = 0
コード例 #6
0
def main():
    args = get_config()

    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(args.n_training_threads)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    # path
    model_dir = Path('./results') / args.env_name / args.algorithm_name / (
        "run" + str(args.seed))
    if args.critic_full_obs:
        run_dir = model_dir / 'adaptive'
    else:
        run_dir = model_dir / 'adaptive_only'
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir))

    print(
        "\n Now we have %i fixed policy! Train Single Adaptive Policy... \n" %
        args.num_policy_candidates)
    args.env_name = args.env_name + "Adaptive"
    policy_candidates = []
    for i in range(args.num_policy_candidates):
        ac = torch.load(
            str(model_dir) + ("/models/Policy%i" % (i + 1)) +
            "-agent0_model.pt")['model'].cpu()
        policy_candidates.append(ac)

    # env
    envs = make_parallel_env(args, policy_candidates)

    #Policy network
    # agent 0
    actor_critic = Policy(envs.observation_space[0],
                          envs.action_space[0],
                          num_agents=args.num_agents,
                          base_kwargs={
                              'lstm': args.lstm,
                              'naive_recurrent': args.naive_recurrent_policy,
                              'recurrent': args.recurrent_policy,
                              'hidden_size': args.hidden_size
                          })

    actor_critic.to(device)
    agent0 = PPO(actor_critic,
                 0,
                 args.clip_param,
                 args.ppo_epoch,
                 args.num_mini_batch,
                 args.data_chunk_length,
                 args.value_loss_coef,
                 args.entropy_coef,
                 logger,
                 lr=args.lr,
                 eps=args.eps,
                 max_grad_norm=args.max_grad_norm,
                 use_clipped_value_loss=args.use_clipped_value_loss)
    #replay buffer
    rollout = RolloutStorage(args.num_agents, 0, args.episode_length,
                             args.n_rollout_threads, envs.observation_space[0],
                             envs.action_space[0],
                             actor_critic.recurrent_hidden_state_size)

    # reset
    if args.critic_full_obs:
        obs, obs_critic, select_opponent = envs.reset()
    else:
        obs, select_opponent = envs.reset()

    # rollout
    if len(envs.observation_space[0]) == 1:
        if args.critic_full_obs:
            rollout.share_obs[0].copy_(
                torch.tensor(obs_critic.reshape(args.n_rollout_threads, -1)))
        else:
            rollout.share_obs[0].copy_(
                torch.tensor(obs.reshape(args.n_rollout_threads, -1)))
        rollout.obs[0].copy_(torch.tensor(obs[:, 0, :]))
        rollout.recurrent_hidden_states.zero_()
        rollout.recurrent_hidden_states_critic.zero_()
        rollout.recurrent_c_states.zero_()
        rollout.recurrent_c_states_critic.zero_()
    else:
        raise NotImplementedError
    rollout.to(device)

    # run
    collective_return = []
    apple_consumption = []
    waste_cleared = []
    sustainability = []
    fire = []

    start = time.time()
    episodes = int(
        args.num_env_steps) // args.episode_length // args.n_rollout_threads
    all_episode = 0
    all_episode_adaptive = np.zeros(args.num_policy_candidates)

    for episode in range(episodes):
        if args.use_linear_lr_decay:
            update_linear_schedule(agent0.optimizer, episode, episodes,
                                   args.lr)

        for step in range(args.episode_length):
            with torch.no_grad():
                value, action0, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic, recurrent_c_states, recurrent_c_states_critic = actor_critic.act(
                    rollout.share_obs[step], rollout.obs[step],
                    rollout.recurrent_hidden_states[step],
                    rollout.recurrent_hidden_states_critic[step],
                    rollout.recurrent_c_states[step],
                    rollout.recurrent_c_states_critic[step],
                    rollout.masks[step])

            # rearrange action
            actions_env = []
            for i in range(args.n_rollout_threads):
                one_hot_action = np.zeros((1, envs.action_space[0].n))
                one_hot_action[0][action0[i]] = 1
                actions_env.append(one_hot_action)

            # Obser reward and next obs
            if args.critic_full_obs:
                obs, obs_critic, select_opponent, reward, done, infos = envs.step(
                    actions_env)
            else:
                obs, select_opponent, reward, done, infos = envs.step(
                    actions_env)

            # If done then clean the history of observations.
            # insert data in buffer
            masks = []
            bad_masks = []
            for i in range(args.num_agents):
                mask = []
                bad_mask = []
                for done_ in done:
                    if done_[i]:
                        mask.append([0.0])
                        bad_mask.append([1.0])
                    else:
                        mask.append([1.0])
                        bad_mask.append([1.0])
                masks.append(torch.FloatTensor(mask))
                bad_masks.append(torch.FloatTensor(bad_mask))

            if len(envs.observation_space[0]) == 1:
                if args.critic_full_obs:
                    rollout.insert(
                        torch.tensor(
                            obs_critic.reshape(args.n_rollout_threads, -1)),
                        torch.tensor(obs[:, 0, :]), recurrent_hidden_states,
                        recurrent_hidden_states_critic, recurrent_c_states,
                        recurrent_c_states_critic, action0,
                        action_log_prob, value,
                        torch.tensor(reward[:, 0].reshape(-1, 1)), masks[0],
                        bad_masks[0])
                else:
                    rollout.insert(
                        torch.tensor(obs.reshape(args.n_rollout_threads, -1)),
                        torch.tensor(obs[:, 0, :]), recurrent_hidden_states,
                        recurrent_hidden_states_critic, recurrent_c_states,
                        recurrent_c_states_critic, action0,
                        action_log_prob, value,
                        torch.tensor(reward[:, 0].reshape(-1, 1)), masks[0],
                        bad_masks[0])
            else:
                raise NotImplementedError

        with torch.no_grad():
            next_value = actor_critic.get_value(
                rollout.share_obs[-1], rollout.obs[-1],
                rollout.recurrent_hidden_states[-1],
                rollout.recurrent_hidden_states_critic[-1],
                rollout.recurrent_c_states[-1],
                rollout.recurrent_c_states_critic[-1],
                rollout.masks[-1]).detach()

        rollout.compute_returns(next_value, args.use_gae, args.gamma,
                                args.gae_lambda, args.use_proper_time_limits)

        # update the network
        value_loss, action_loss, dist_entropy = agent0.update(rollout)

        if args.env_name == "StagHuntAdaptive":
            coop_num = []
            defect_num = []
            coopdefect_num = []
            defectcoop_num = []
            rew = []
            for info in infos:
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'defect&defect_num' in info.keys():
                    defect_num.append(info['defect&defect_num'])
                if 'coop&defect_num' in info.keys():
                    coopdefect_num.append(info['coop&defect_num'])
                if 'defect&coop_num' in info.keys():
                    defectcoop_num.append(info['defect&coop_num'])
            for i in range(args.n_rollout_threads):
                rew.append(rollout.rewards[:, i, :].sum().cpu().numpy())

            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/defect&defect_num_per_episode',
                    {'defect&defect_num_per_episode': defect_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/coop&defect_num_per_episode',
                    {'coop&defect_num_per_episode': coopdefect_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/defect&coop_num_per_episode',
                    {'defect&coop_num_per_episode': defectcoop_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) + '/reward',
                    {'reward': np.mean(np.array(rew[i]))},
                    all_episode_adaptive[select_opponent[i]])
                all_episode_adaptive[select_opponent[i]] += 1
        elif args.env_name == "StagHuntGWAdaptive":
            collective_return = []
            coop_num = []
            gore1_num = []
            gore2_num = []
            hare1_num = []
            hare2_num = []
            for info in infos:
                if 'collective_return' in info.keys():
                    collective_return.append(info['collective_return'])
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'gore1_num' in info.keys():
                    gore1_num.append(info['gore1_num'])
                if 'gore2_num' in info.keys():
                    gore2_num.append(info['gore2_num'])
                if 'hare1_num' in info.keys():
                    hare1_num.append(info['hare1_num'])
                if 'hare2_num' in info.keys():
                    hare2_num.append(info['hare2_num'])

            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/collective_return',
                    {'collective_return': collective_return[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/gore1_num_per_episode',
                    {'gore1_num_per_episode': gore1_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/gore2_num_per_episode',
                    {'gore2_num_per_episode': gore2_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/hare1_num_per_episode',
                    {'hare1_num_per_episode': hare1_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                logger.add_scalars(
                    'Policy-' + str(select_opponent[i] + 1) +
                    '/hare2_num_per_episode',
                    {'hare2_num_per_episode': hare2_num[i]},
                    all_episode_adaptive[select_opponent[i]])
                all_episode_adaptive[select_opponent[i]] += 1

        if args.critic_full_obs:
            obs, obs_critic, select_opponent = envs.reset()
        else:
            obs, select_opponent = envs.reset()

        if len(envs.observation_space[0]) == 1:
            if args.critic_full_obs:
                rollout.share_obs[0].copy_(
                    torch.tensor(obs_critic.reshape(args.n_rollout_threads,
                                                    -1)))
            else:
                rollout.share_obs[0].copy_(
                    torch.tensor(obs.reshape(args.n_rollout_threads, -1)))
            rollout.obs[0].copy_(torch.tensor(obs[:, 0, :]))
            rollout.recurrent_hidden_states.zero_()
            rollout.recurrent_hidden_states_critic.zero_()
            rollout.recurrent_c_states.zero_()
            rollout.recurrent_c_states_critic.zero_()
            rollout.masks[0].copy_(torch.ones(args.n_rollout_threads, 1))
            rollout.bad_masks[0].copy_(torch.ones(args.n_rollout_threads, 1))
        else:
            raise NotImplementedError
        rollout.to(device)

        if (episode % args.save_interval == 0 or episode == episodes - 1):
            torch.save({'model': actor_critic},
                       str(save_dir) + "/agent0_model.pt")

        # log information
        if episode % args.log_interval == 0:
            total_num_steps = (
                episode + 1) * args.episode_length * args.n_rollout_threads
            end = time.time()
            print(
                "\n Updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(episode, episodes, total_num_steps, args.num_env_steps,
                        int(total_num_steps / (end - start))))
            print("value loss: agent0--" + str(value_loss))
    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()
コード例 #7
0
def main():
    args = get_config()
    run = wandb.init(project='curriculum',name=str(args.algorithm_name) + "_seed" + str(args.seed))
    # run = wandb.init(project='check',name='separate_reward')
    
    assert (args.share_policy == True and args.scenario_name == 'simple_speaker_listener') == False, ("The simple_speaker_listener scenario can not use shared policy. Please check the config.py.")

    # seed
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    
    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(1)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)
    
    # path
    model_dir = Path('./results') / args.env_name / args.scenario_name / args.algorithm_name
    node_dir = Path('./node') / args.env_name / args.scenario_name / args.algorithm_name
    if not model_dir.exists():
        curr_run = 'run1'
    else:
        exst_run_nums = [int(str(folder.name).split('run')[1]) for folder in model_dir.iterdir() if str(folder.name).startswith('run')]
        if len(exst_run_nums) == 0:
            curr_run = 'run1'
        else:
            curr_run = 'run%i' % (max(exst_run_nums) + 1)
    if not node_dir.exists():
        node_curr_run = 'run1'
    else:
        exst_run_nums = [int(str(folder.name).split('run')[1]) for folder in node_dir.iterdir() if str(folder.name).startswith('run')]
        if len(exst_run_nums) == 0:
            node_curr_run = 'run1'
        else:
            node_curr_run = 'run%i' % (max(exst_run_nums) + 1)

    run_dir = model_dir / curr_run
    save_node_dir = node_dir / node_curr_run
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir)) 

    # env
    envs = make_parallel_env(args)
    num_agents = args.num_agents
    #Policy network
    if args.share_policy:
        actor_base = ATTBase_actor_dist_add(envs.observation_space[0].shape[0], envs.action_space[0], num_agents)
        critic_base = ATTBase_critic_add(envs.observation_space[0].shape[0], num_agents)
        actor_critic = Policy3(envs.observation_space[0], 
                    envs.action_space[0],
                    num_agents = num_agents,
                    base=None,
                    actor_base=actor_base,
                    critic_base=critic_base,
                    base_kwargs={'naive_recurrent': args.naive_recurrent_policy,
                                 'recurrent': args.recurrent_policy,
                                 'hidden_size': args.hidden_size,
                                 'attn': args.attn,                                 
                                 'attn_size': args.attn_size,
                                 'attn_N': args.attn_N,
                                 'attn_heads': args.attn_heads,
                                 'dropout': args.dropout,
                                 'use_average_pool': args.use_average_pool,
                                 'use_common_layer':args.use_common_layer,
                                 'use_feature_normlization':args.use_feature_normlization,
                                 'use_feature_popart':args.use_feature_popart,
                                 'use_orthogonal':args.use_orthogonal,
                                 'layer_N':args.layer_N,
                                 'use_ReLU':args.use_ReLU
                                 },
                    device = device)
        actor_critic.to(device)
        # algorithm
        agents = PPO3(actor_critic,
                   args.clip_param,
                   args.ppo_epoch,
                   args.num_mini_batch,
                   args.data_chunk_length,
                   args.value_loss_coef,
                   args.entropy_coef,
                   logger,
                   lr=args.lr,
                   eps=args.eps,
                   weight_decay=args.weight_decay,
                   max_grad_norm=args.max_grad_norm,
                   use_max_grad_norm=args.use_max_grad_norm,
                   use_clipped_value_loss= args.use_clipped_value_loss,
                   use_common_layer=args.use_common_layer,
                   use_huber_loss=args.use_huber_loss,
                   huber_delta=args.huber_delta,
                   use_popart=args.use_popart,
                   device=device)
                   
        # #replay buffer
        # rollouts = RolloutStorage(num_agents,
        #             args.episode_length, 
        #             args.n_rollout_threads,
        #             envs.observation_space[0], 
        #             envs.action_space[0],
        #             args.hidden_size)        
    else:
        actor_critic = []
        agents = []
        rollouts = []
        for agent_id in range(num_agents):
            ac = Policy(envs.observation_space, 
                      envs.action_space[agent_id],
                      num_agents = agent_id, # here is special
                      base_kwargs={'naive_recurrent': args.naive_recurrent_policy,
                                 'recurrent': args.recurrent_policy,
                                 'hidden_size': args.hidden_size,
                                 'attn': args.attn,                                 
                                 'attn_size': args.attn_size,
                                 'attn_N': args.attn_N,
                                 'attn_heads': args.attn_heads,
                                 'dropout': args.dropout,
                                 'use_average_pool': args.use_average_pool,
                                 'use_common_layer':args.use_common_layer,
                                 'use_feature_normlization':args.use_feature_normlization,
                                 'use_feature_popart':args.use_feature_popart,
                                 'use_orthogonal':args.use_orthogonal,
                                 'layer_N':args.layer_N,
                                 'use_ReLU':args.use_ReLU
                                 },
                      device = device)
            ac.to(device)
            # algorithm
            agent = PPO(ac,
                   args.clip_param,
                   args.ppo_epoch,
                   args.num_mini_batch,
                   args.data_chunk_length,
                   args.value_loss_coef,
                   args.entropy_coef,
                   logger,
                   lr=args.lr,
                   eps=args.eps,
                   weight_decay=args.weight_decay,
                   max_grad_norm=args.max_grad_norm,
                   use_max_grad_norm=args.use_max_grad_norm,
                   use_clipped_value_loss= args.use_clipped_value_loss,
                   use_common_layer=args.use_common_layer,
                   use_huber_loss=args.use_huber_loss,
                   huber_delta=args.huber_delta,
                   use_popart=args.use_popart,
                   device=device)
                               
            actor_critic.append(ac)
            agents.append(agent) 
              
            #replay buffer
            ro = SingleRolloutStorage(agent_id,
                    args.episode_length, 
                    args.n_rollout_threads,
                    envs.observation_space, 
                    envs.action_space,
                    args.hidden_size)
            rollouts.append(ro)
    
   
    boundary = 3
    start_boundary = [-0.3,0.3,-0.3,0.3] # 分别代表x的范围和y的范围
    # start_boundary = [2.4,3.0,2.4,3.0]
    max_step = 0.6
    N_easy = 0
    test_flag = 0
    reproduce_flag = 0
    target_num = 4
    last_agent_num = 4
    now_agent_num = num_agents
    mean_cover_rate = 0
    eval_frequency = 2 #需要fix几个回合
    check_frequency = 1
    save_node_frequency = 5
    save_node_flag = True
    save_90_flag = True
    historical_length = 5
    random.seed(args.seed)
    np.random.seed(args.seed)


    # init the Gan
    gan_configs['goal_range'] = boundary
    gan_configs['goal_center'] = np.zeros((num_agents + num_agents)* 2, dtype=float)
    gan_configs['goal_size'] = (num_agents + num_agents)*2
    gan = StateGAN(gan_configs = gan_configs, state_range=gan_configs['goal_range'])
    feasible_goals = generate_initial_goals(num_case = 10000, start_boundary = start_boundary, agent_num = args.num_agents)                            
    dis_loss, gen_loss = gan.pretrain(states=feasible_goals, outer_iters=gan_configs['gan_outer_iters'])
    print('discriminator_loss:',str(dis_loss.cpu()), 'generator_loss:',str(gen_loss.cpu()))
    
    # init the StateCollection
    all_goals = StateCollection(distance_threshold=goal_configs['coll_eps'])

    # run
    begin = time.time()
    episodes = int(args.num_env_steps) // args.episode_length // args.n_rollout_threads // eval_frequency
    curriculum_episode = 0
    current_timestep = 0
    one_length = args.n_rollout_threads
    starts_length = args.n_rollout_threads
    num_envs = args.n_rollout_threads

    for episode in range(episodes):
        if args.use_linear_lr_decay:# decrease learning rate linearly
            if args.share_policy:   
                update_linear_schedule(agents.optimizer, episode, episodes, args.lr)  
            else:     
                for agent_id in range(num_agents):
                    update_linear_schedule(agents[agent_id].optimizer, episode, episodes, args.lr)           



        raw_goals, _ = gan.sample_states_with_noise(goal_configs['num_new_goals'])
        # replay buffer
        if all_goals.size > 0:
            old_goals = all_goals.sample(goal_configs['num_old_goals'])
            goals = np.vstack([raw_goals, old_goals])
        else:
            goals = raw_goals   
        if goals.shape[0] < num_envs:
            add_num = num_envs - goals.shape[0]
            goals = np.vstack([goals, goals[:add_num]]) #补齐到num_new_goals+num_old_goals   
        # generate the starts
        starts = numpy_to_list(goals, list_length=num_envs, shape=(num_agents*2,2))

        for times in range(eval_frequency):
            obs = envs.new_starts_obs(starts, now_agent_num, starts_length)
            #replay buffer
            rollouts = RolloutStorage(num_agents,
                        args.episode_length, 
                        starts_length,
                        envs.observation_space[0], 
                        envs.action_space[0],
                        args.hidden_size) 
            # replay buffer init
            if args.share_policy: 
                share_obs = obs.reshape(starts_length, -1)        
                share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)    
                rollouts.share_obs[0] = share_obs.copy() 
                rollouts.obs[0] = obs.copy()               
                rollouts.recurrent_hidden_states = np.zeros(rollouts.recurrent_hidden_states.shape).astype(np.float32)
                rollouts.recurrent_hidden_states_critic = np.zeros(rollouts.recurrent_hidden_states_critic.shape).astype(np.float32)
            else:
                share_obs = []
                for o in obs:
                    share_obs.append(list(itertools.chain(*o)))
                share_obs = np.array(share_obs)
                for agent_id in range(num_agents):    
                    rollouts[agent_id].share_obs[0] = share_obs.copy()
                    rollouts[agent_id].obs[0] = np.array(list(obs[:,agent_id])).copy()               
                    rollouts[agent_id].recurrent_hidden_states = np.zeros(rollouts[agent_id].recurrent_hidden_states.shape).astype(np.float32)
                    rollouts[agent_id].recurrent_hidden_states_critic = np.zeros(rollouts[agent_id].recurrent_hidden_states_critic.shape).astype(np.float32)
            step_cover_rate = np.zeros(shape=(one_length,args.episode_length))
            for step in range(args.episode_length):
                # Sample actions
                values = []
                actions= []
                action_log_probs = []
                recurrent_hidden_statess = []
                recurrent_hidden_statess_critic = []
                
                with torch.no_grad():                
                    for agent_id in range(num_agents):
                        if args.share_policy:
                            actor_critic.eval()
                            value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(agent_id,
                                torch.FloatTensor(rollouts.share_obs[step,:,agent_id]), 
                                torch.FloatTensor(rollouts.obs[step,:,agent_id]), 
                                torch.FloatTensor(rollouts.recurrent_hidden_states[step,:,agent_id]), 
                                torch.FloatTensor(rollouts.recurrent_hidden_states_critic[step,:,agent_id]),
                                torch.FloatTensor(rollouts.masks[step,:,agent_id]))
                        else:
                            actor_critic[agent_id].eval()
                            value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic[agent_id].act(agent_id,
                                torch.FloatTensor(rollouts[agent_id].share_obs[step,:]), 
                                torch.FloatTensor(rollouts[agent_id].obs[step,:]), 
                                torch.FloatTensor(rollouts[agent_id].recurrent_hidden_states[step,:]), 
                                torch.FloatTensor(rollouts[agent_id].recurrent_hidden_states_critic[step,:]),
                                torch.FloatTensor(rollouts[agent_id].masks[step,:]))
                            
                        values.append(value.detach().cpu().numpy())
                        actions.append(action.detach().cpu().numpy())
                        action_log_probs.append(action_log_prob.detach().cpu().numpy())
                        recurrent_hidden_statess.append(recurrent_hidden_states.detach().cpu().numpy())
                        recurrent_hidden_statess_critic.append(recurrent_hidden_states_critic.detach().cpu().numpy())
                # rearrange action
                actions_env = []
                for i in range(starts_length):
                    one_hot_action_env = []
                    for agent_id in range(num_agents):
                        if envs.action_space[agent_id].__class__.__name__ == 'MultiDiscrete':
                            uc_action = []
                            for j in range(envs.action_space[agent_id].shape):
                                uc_one_hot_action = np.zeros(envs.action_space[agent_id].high[j]+1)
                                uc_one_hot_action[actions[agent_id][i][j]] = 1
                                uc_action.append(uc_one_hot_action)
                            uc_action = np.concatenate(uc_action)
                            one_hot_action_env.append(uc_action)
                                
                        elif envs.action_space[agent_id].__class__.__name__ == 'Discrete':    
                            one_hot_action = np.zeros(envs.action_space[agent_id].n)
                            one_hot_action[actions[agent_id][i]] = 1
                            one_hot_action_env.append(one_hot_action)
                        else:
                            raise NotImplementedError
                    actions_env.append(one_hot_action_env)
                
                # Obser reward and next obs
                obs, rewards, dones, infos, _ = envs.step(actions_env, starts_length, num_agents)
                cover_rate_list = []
                for env_id in range(one_length):
                    cover_rate_list.append(infos[env_id][0]['cover_rate'])
                step_cover_rate[:,step] = np.array(cover_rate_list)
                # step_cover_rate[:,step] = np.array(infos)[0:one_length,0]

                # If done then clean the history of observations.
                # insert data in buffer
                masks = []
                for i, done in enumerate(dones): 
                    mask = []               
                    for agent_id in range(num_agents): 
                        if done[agent_id]:    
                            recurrent_hidden_statess[agent_id][i] = np.zeros(args.hidden_size).astype(np.float32)
                            recurrent_hidden_statess_critic[agent_id][i] = np.zeros(args.hidden_size).astype(np.float32)    
                            mask.append([0.0])
                        else:
                            mask.append([1.0])
                    masks.append(mask)
                                
                if args.share_policy: 
                    share_obs = obs.reshape(starts_length, -1)        
                    share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)    
                    rollouts.insert(share_obs, 
                                obs, 
                                np.array(recurrent_hidden_statess).transpose(1,0,2), 
                                np.array(recurrent_hidden_statess_critic).transpose(1,0,2), 
                                np.array(actions).transpose(1,0,2),
                                np.array(action_log_probs).transpose(1,0,2), 
                                np.array(values).transpose(1,0,2),
                                rewards, 
                                masks)
                else:
                    share_obs = []
                    for o in obs:
                        share_obs.append(list(itertools.chain(*o)))
                    share_obs = np.array(share_obs)
                    for agent_id in range(num_agents):
                        rollouts[agent_id].insert(share_obs, 
                                np.array(list(obs[:,agent_id])), 
                                np.array(recurrent_hidden_statess[agent_id]), 
                                np.array(recurrent_hidden_statess_critic[agent_id]), 
                                np.array(actions[agent_id]),
                                np.array(action_log_probs[agent_id]), 
                                np.array(values[agent_id]),
                                rewards[:,agent_id], 
                                np.array(masks)[:,agent_id])
            # logger.add_scalars('agent/training_cover_rate',{'training_cover_rate': np.mean(np.mean(step_cover_rate[:,-historical_length:],axis=1))}, current_timestep)
            wandb.log({'training_cover_rate': np.mean(np.mean(step_cover_rate[:,-historical_length:],axis=1))}, current_timestep)
            print('training_cover_rate: ', np.mean(np.mean(step_cover_rate[:,-historical_length:],axis=1)))
            current_timestep += args.episode_length * starts_length
            curriculum_episode += 1
            
            #region train the gan

            
            if times == 1:
                start_time = time.time()
                filtered_raw_goals = []
                labels = np.zeros((num_envs, 1), dtype = int)
                for i in range(num_envs):
                    R_i = np.mean(step_cover_rate[i, -goal_configs['historical_length']:])
                    if R_i < goal_configs['R_max'] and R_i > goal_configs['R_min']:
                        labels[i] = 1
                        filtered_raw_goals.append(goals[i])
                gan.train(goals, labels)
                all_goals.append(filtered_raw_goals)
                end_time = time.time()
                print("Gan training time: %.2f"%(end_time-start_time))

            with torch.no_grad():  # get value and com
                for agent_id in range(num_agents):         
                    if args.share_policy: 
                        actor_critic.eval()                
                        next_value,_,_ = actor_critic.get_value(agent_id,
                                                    torch.FloatTensor(rollouts.share_obs[-1,:,agent_id]), 
                                                    torch.FloatTensor(rollouts.obs[-1,:,agent_id]), 
                                                    torch.FloatTensor(rollouts.recurrent_hidden_states[-1,:,agent_id]),
                                                    torch.FloatTensor(rollouts.recurrent_hidden_states_critic[-1,:,agent_id]),
                                                    torch.FloatTensor(rollouts.masks[-1,:,agent_id]))
                        next_value = next_value.detach().cpu().numpy()
                        rollouts.compute_returns(agent_id,
                                        next_value, 
                                        args.use_gae, 
                                        args.gamma,
                                        args.gae_lambda, 
                                        args.use_proper_time_limits,
                                        args.use_popart,
                                        agents.value_normalizer)
                    else:
                        actor_critic[agent_id].eval()
                        next_value,_,_ = actor_critic[agent_id].get_value(agent_id,
                                                                torch.FloatTensor(rollouts[agent_id].share_obs[-1,:]), 
                                                                torch.FloatTensor(rollouts[agent_id].obs[-1,:]), 
                                                                torch.FloatTensor(rollouts[agent_id].recurrent_hidden_states[-1,:]),
                                                                torch.FloatTensor(rollouts[agent_id].recurrent_hidden_states_critic[-1,:]),
                                                                torch.FloatTensor(rollouts[agent_id].masks[-1,:]))
                        next_value = next_value.detach().cpu().numpy()
                        rollouts[agent_id].compute_returns(next_value, 
                                                args.use_gae, 
                                                args.gamma,
                                                args.gae_lambda, 
                                                args.use_proper_time_limits,
                                                args.use_popart,
                                                agents[agent_id].value_normalizer)

            # update the network
            if args.share_policy:
                actor_critic.train()
                value_loss, action_loss, dist_entropy = agents.update_share_asynchronous(last_agent_num, rollouts, False, initial_optimizer=False) 
                print('value_loss: ', value_loss)
                wandb.log(
                    {'value_loss': value_loss},
                    current_timestep)
                rew = []
                for i in range(rollouts.rewards.shape[1]):
                    rew.append(np.sum(rollouts.rewards[:,i]))
                wandb.log(
                    {'average_episode_reward': np.mean(rew)},
                    current_timestep)
                # clean the buffer and reset
                rollouts.after_update()
            else:
                value_losses = []
                action_losses = []
                dist_entropies = [] 
                
                for agent_id in range(num_agents):
                    actor_critic[agent_id].train()
                    value_loss, action_loss, dist_entropy = agents[agent_id].update_single(agent_id, rollouts[agent_id])
                    value_losses.append(value_loss)
                    action_losses.append(action_loss)
                    dist_entropies.append(dist_entropy)
                        
                    rew = []
                    for i in range(rollouts[agent_id].rewards.shape[1]):
                        rew.append(np.sum(rollouts[agent_id].rewards[:,i]))

                    logger.add_scalars('agent%i/average_episode_reward'%agent_id,
                        {'average_episode_reward': np.mean(rew)},
                        (episode+1) * args.episode_length * one_length*eval_frequency)
                    
                    rollouts[agent_id].after_update()


        # test
        if episode % check_frequency==0:
            obs, _ = envs.reset(num_agents)
            episode_length = 70
            #replay buffer
            rollouts = RolloutStorage(num_agents,
                        episode_length, 
                        args.n_rollout_threads,
                        envs.observation_space[0], 
                        envs.action_space[0],
                        args.hidden_size) 
            # replay buffer init
            if args.share_policy: 
                share_obs = obs.reshape(args.n_rollout_threads, -1)        
                share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)    
                rollouts.share_obs[0] = share_obs.copy() 
                rollouts.obs[0] = obs.copy()               
                rollouts.recurrent_hidden_states = np.zeros(rollouts.recurrent_hidden_states.shape).astype(np.float32)
                rollouts.recurrent_hidden_states_critic = np.zeros(rollouts.recurrent_hidden_states_critic.shape).astype(np.float32)
            else:
                share_obs = []
                for o in obs:
                    share_obs.append(list(itertools.chain(*o)))
                share_obs = np.array(share_obs)
                for agent_id in range(num_agents):    
                    rollouts[agent_id].share_obs[0] = share_obs.copy()
                    rollouts[agent_id].obs[0] = np.array(list(obs[:,agent_id])).copy()               
                    rollouts[agent_id].recurrent_hidden_states = np.zeros(rollouts[agent_id].recurrent_hidden_states.shape).astype(np.float32)
                    rollouts[agent_id].recurrent_hidden_states_critic = np.zeros(rollouts[agent_id].recurrent_hidden_states_critic.shape).astype(np.float32)
            test_cover_rate = np.zeros(shape=(args.n_rollout_threads,episode_length))
            for step in range(episode_length):
                # Sample actions
                values = []
                actions= []
                action_log_probs = []
                recurrent_hidden_statess = []
                recurrent_hidden_statess_critic = []
                
                with torch.no_grad():                
                    for agent_id in range(num_agents):
                        if args.share_policy:
                            actor_critic.eval()
                            value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(agent_id,
                                torch.FloatTensor(rollouts.share_obs[step,:,agent_id]), 
                                torch.FloatTensor(rollouts.obs[step,:,agent_id]), 
                                torch.FloatTensor(rollouts.recurrent_hidden_states[step,:,agent_id]), 
                                torch.FloatTensor(rollouts.recurrent_hidden_states_critic[step,:,agent_id]),
                                torch.FloatTensor(rollouts.masks[step,:,agent_id]),deterministic=True)
                        else:
                            actor_critic[agent_id].eval()
                            value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic[agent_id].act(agent_id,
                                torch.FloatTensor(rollouts[agent_id].share_obs[step,:]), 
                                torch.FloatTensor(rollouts[agent_id].obs[step,:]), 
                                torch.FloatTensor(rollouts[agent_id].recurrent_hidden_states[step,:]), 
                                torch.FloatTensor(rollouts[agent_id].recurrent_hidden_states_critic[step,:]),
                                torch.FloatTensor(rollouts[agent_id].masks[step,:]),deterministic=True)
                            
                        values.append(value.detach().cpu().numpy())
                        actions.append(action.detach().cpu().numpy())
                        action_log_probs.append(action_log_prob.detach().cpu().numpy())
                        recurrent_hidden_statess.append(recurrent_hidden_states.detach().cpu().numpy())
                        recurrent_hidden_statess_critic.append(recurrent_hidden_states_critic.detach().cpu().numpy())
                
                # rearrange action
                actions_env = []
                for i in range(args.n_rollout_threads):
                    one_hot_action_env = []
                    for agent_id in range(num_agents):
                        if envs.action_space[agent_id].__class__.__name__ == 'MultiDiscrete':
                            uc_action = []
                            for j in range(envs.action_space[agent_id].shape):
                                uc_one_hot_action = np.zeros(envs.action_space[agent_id].high[j]+1)
                                uc_one_hot_action[actions[agent_id][i][j]] = 1
                                uc_action.append(uc_one_hot_action)
                            uc_action = np.concatenate(uc_action)
                            one_hot_action_env.append(uc_action)
                                
                        elif envs.action_space[agent_id].__class__.__name__ == 'Discrete':    
                            one_hot_action = np.zeros(envs.action_space[agent_id].n)
                            one_hot_action[actions[agent_id][i]] = 1
                            one_hot_action_env.append(one_hot_action)
                        else:
                            raise NotImplementedError
                    actions_env.append(one_hot_action_env)
                
                # Obser reward and next obs
                obs, rewards, dones, infos, _ = envs.step(actions_env, args.n_rollout_threads, num_agents)
                cover_rate_list = []
                for env_id in range(args.n_rollout_threads):
                    cover_rate_list.append(infos[env_id][0]['cover_rate'])
                test_cover_rate[:,step] = np.array(cover_rate_list)
                # test_cover_rate[:,step] = np.array(infos)[:,0]

                # If done then clean the history of observations.
                # insert data in buffer
                masks = []
                for i, done in enumerate(dones): 
                    mask = []               
                    for agent_id in range(num_agents): 
                        if done[agent_id]:    
                            recurrent_hidden_statess[agent_id][i] = np.zeros(args.hidden_size).astype(np.float32)
                            recurrent_hidden_statess_critic[agent_id][i] = np.zeros(args.hidden_size).astype(np.float32)    
                            mask.append([0.0])
                        else:
                            mask.append([1.0])
                    masks.append(mask)
                                
                if args.share_policy: 
                    share_obs = obs.reshape(args.n_rollout_threads, -1)        
                    share_obs = np.expand_dims(share_obs,1).repeat(num_agents,axis=1)    
                    
                    rollouts.insert(share_obs, 
                                obs, 
                                np.array(recurrent_hidden_statess).transpose(1,0,2), 
                                np.array(recurrent_hidden_statess_critic).transpose(1,0,2), 
                                np.array(actions).transpose(1,0,2),
                                np.array(action_log_probs).transpose(1,0,2), 
                                np.array(values).transpose(1,0,2),
                                rewards, 
                                masks)
                else:
                    share_obs = []
                    for o in obs:
                        share_obs.append(list(itertools.chain(*o)))
                    share_obs = np.array(share_obs)
                    for agent_id in range(num_agents):
                        rollouts[agent_id].insert(share_obs, 
                                np.array(list(obs[:,agent_id])), 
                                np.array(recurrent_hidden_statess[agent_id]), 
                                np.array(recurrent_hidden_statess_critic[agent_id]), 
                                np.array(actions[agent_id]),
                                np.array(action_log_probs[agent_id]), 
                                np.array(values[agent_id]),
                                rewards[:,agent_id], 
                                np.array(masks)[:,agent_id])

            # logger.add_scalars('agent/cover_rate_1step',{'cover_rate_1step': np.mean(test_cover_rate[:,-1])},current_timestep)
            # logger.add_scalars('agent/cover_rate_5step',{'cover_rate_5step': np.mean(np.mean(test_cover_rate[:,-historical_length:],axis=1))}, current_timestep)
            rew = []
            for i in range(rollouts.rewards.shape[1]):
                rew.append(np.sum(rollouts.rewards[:,i]))
            wandb.log(
                {'eval_episode_reward': np.mean(rew)},
                current_timestep)
            wandb.log({'cover_rate_1step': np.mean(test_cover_rate[:,-1])},current_timestep)
            wandb.log({'cover_rate_5step': np.mean(np.mean(test_cover_rate[:,-historical_length:],axis=1))}, current_timestep)
            mean_cover_rate = np.mean(np.mean(test_cover_rate[:,-historical_length:],axis=1))
            if mean_cover_rate >= 0.9 and args.algorithm_name=='ours' and save_90_flag:
                torch.save({'model': actor_critic}, str(save_dir) + "/cover09_agent_model.pt")
                save_90_flag = False
            print('test_agent_num: ', last_agent_num)
            print('test_mean_cover_rate: ', mean_cover_rate)

        total_num_steps = current_timestep

        if (episode % args.save_interval == 0 or episode == episodes - 1):# save for every interval-th episode or for the last epoch
            if args.share_policy:
                torch.save({
                            'model': actor_critic
                            }, 
                            str(save_dir) + "/agent_model.pt")
            else:
                for agent_id in range(num_agents):                                                  
                    torch.save({
                                'model': actor_critic[agent_id]
                                }, 
                                str(save_dir) + "/agent%i_model" % agent_id + ".pt")

        # log information
        if episode % args.log_interval == 0:
            end = time.time()
            print("\n Scenario {} Algo {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(args.scenario_name,
                        args.algorithm_name,
                        episode, 
                        episodes,
                        total_num_steps,
                        args.num_env_steps,
                        int(total_num_steps / (end - begin))))
            if args.share_policy:
                print("value loss of agent: " + str(value_loss))
            else:
                for agent_id in range(num_agents):
                    print("value loss of agent%i: " % agent_id + str(value_losses[agent_id]))
                
    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()
    envs.close()
コード例 #8
0
def main():
    args = get_config()

    # seed
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    np.random.seed(args.seed)

    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(args.n_training_threads)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)

    # path
    model_dir = Path(
        './results') / args.env_name / args.map_name / args.algorithm_name
    if not model_dir.exists():
        curr_run = 'run1'
    else:
        exst_run_nums = [
            int(str(folder.name).split('run')[1])
            for folder in model_dir.iterdir()
            if str(folder.name).startswith('run')
        ]
        if len(exst_run_nums) == 0:
            curr_run = 'run1'
        else:
            curr_run = 'run%i' % (max(exst_run_nums) + 1)

    run_dir = model_dir / curr_run
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir))

    # env
    envs = make_parallel_env(args)
    if args.eval:
        eval_env = make_eval_env(args)
    num_agents = get_map_params(args.map_name)["n_agents"]
    #Policy network

    if args.share_policy:
        actor_critic = Policy(envs.observation_space[0],
                              envs.action_space[0],
                              num_agents=num_agents,
                              gain=args.gain,
                              base_kwargs={
                                  'naive_recurrent':
                                  args.naive_recurrent_policy,
                                  'recurrent': args.recurrent_policy,
                                  'hidden_size': args.hidden_size,
                                  'recurrent_N': args.recurrent_N,
                                  'attn': args.attn,
                                  'attn_only_critic': args.attn_only_critic,
                                  'attn_size': args.attn_size,
                                  'attn_N': args.attn_N,
                                  'attn_heads': args.attn_heads,
                                  'dropout': args.dropout,
                                  'use_average_pool': args.use_average_pool,
                                  'use_common_layer': args.use_common_layer,
                                  'use_feature_normlization':
                                  args.use_feature_normlization,
                                  'use_feature_popart':
                                  args.use_feature_popart,
                                  'use_orthogonal': args.use_orthogonal,
                                  'layer_N': args.layer_N,
                                  'use_ReLU': args.use_ReLU,
                                  'use_same_dim': args.use_same_dim
                              },
                              device=device)
        actor_critic.to(device)
        # algorithm
        agents = PPO(actor_critic,
                     args.clip_param,
                     args.ppo_epoch,
                     args.num_mini_batch,
                     args.data_chunk_length,
                     args.value_loss_coef,
                     args.entropy_coef,
                     logger,
                     lr=args.lr,
                     eps=args.eps,
                     weight_decay=args.weight_decay,
                     max_grad_norm=args.max_grad_norm,
                     use_max_grad_norm=args.use_max_grad_norm,
                     use_clipped_value_loss=args.use_clipped_value_loss,
                     use_common_layer=args.use_common_layer,
                     use_huber_loss=args.use_huber_loss,
                     huber_delta=args.huber_delta,
                     use_popart=args.use_popart,
                     use_value_high_masks=args.use_value_high_masks,
                     device=device)

        #replay buffer
        rollouts = RolloutStorage(num_agents, args.episode_length,
                                  args.n_rollout_threads,
                                  envs.observation_space[0],
                                  envs.action_space[0], args.hidden_size)
    else:
        actor_critic = []
        agents = []
        for agent_id in range(num_agents):
            ac = Policy(envs.observation_space[0],
                        envs.action_space[0],
                        num_agents=num_agents,
                        gain=args.gain,
                        base_kwargs={
                            'naive_recurrent': args.naive_recurrent_policy,
                            'recurrent': args.recurrent_policy,
                            'hidden_size': args.hidden_size,
                            'recurrent_N': args.recurrent_N,
                            'attn': args.attn,
                            'attn_only_critic': args.attn_only_critic,
                            'attn_size': args.attn_size,
                            'attn_N': args.attn_N,
                            'attn_heads': args.attn_heads,
                            'dropout': args.dropout,
                            'use_average_pool': args.use_average_pool,
                            'use_common_layer': args.use_common_layer,
                            'use_feature_normlization':
                            args.use_feature_normlization,
                            'use_feature_popart': args.use_feature_popart,
                            'use_orthogonal': args.use_orthogonal,
                            'layer_N': args.layer_N,
                            'use_ReLU': args.use_ReLU,
                            'use_same_dim': args.use_same_dim
                        },
                        device=device)
            ac.to(device)
            # algorithm
            agent = PPO(ac,
                        args.clip_param,
                        args.ppo_epoch,
                        args.num_mini_batch,
                        args.data_chunk_length,
                        args.value_loss_coef,
                        args.entropy_coef,
                        logger,
                        lr=args.lr,
                        eps=args.eps,
                        weight_decay=args.weight_decay,
                        max_grad_norm=args.max_grad_norm,
                        use_max_grad_norm=args.use_max_grad_norm,
                        use_clipped_value_loss=args.use_clipped_value_loss,
                        use_common_layer=args.use_common_layer,
                        use_huber_loss=args.use_huber_loss,
                        huber_delta=args.huber_delta,
                        use_popart=args.use_popart,
                        use_value_high_masks=args.use_value_high_masks,
                        device=device)

            actor_critic.append(ac)
            agents.append(agent)

        #replay buffer
        rollouts = RolloutStorage(num_agents, args.episode_length,
                                  args.n_rollout_threads,
                                  envs.observation_space[0],
                                  envs.action_space[0], args.hidden_size)

    # reset env
    obs, available_actions = envs.reset()

    # replay buffer
    if len(envs.observation_space[0]) == 3:
        share_obs = obs.reshape(args.n_rollout_threads, -1,
                                envs.observation_space[0][1],
                                envs.observation_space[0][2])
    else:
        share_obs = obs.reshape(args.n_rollout_threads, -1)

    share_obs = np.expand_dims(share_obs, 1).repeat(num_agents, axis=1)
    rollouts.share_obs[0] = share_obs.copy()
    rollouts.obs[0] = obs.copy()
    rollouts.available_actions[0] = available_actions.copy()
    rollouts.recurrent_hidden_states = np.zeros(
        rollouts.recurrent_hidden_states.shape).astype(np.float32)
    rollouts.recurrent_hidden_states_critic = np.zeros(
        rollouts.recurrent_hidden_states_critic.shape).astype(np.float32)

    # run
    start = time.time()
    episodes = int(
        args.num_env_steps) // args.episode_length // args.n_rollout_threads
    timesteps = 0
    last_battles_game = np.zeros(args.n_rollout_threads)
    last_battles_won = np.zeros(args.n_rollout_threads)

    for episode in range(episodes):
        if args.use_linear_lr_decay:  # decrease learning rate linearly
            if args.share_policy:
                update_linear_schedule(agents.optimizer, episode, episodes,
                                       args.lr)
            else:
                for agent_id in range(num_agents):
                    update_linear_schedule(agents[agent_id].optimizer, episode,
                                           episodes, args.lr)

        for step in range(args.episode_length):
            # Sample actions
            values = []
            actions = []
            action_log_probs = []
            recurrent_hidden_statess = []
            recurrent_hidden_statess_critic = []

            with torch.no_grad():
                for agent_id in range(num_agents):
                    if args.share_policy:
                        actor_critic.eval()
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(
                            agent_id,
                            torch.tensor(rollouts.share_obs[step, :,
                                                            agent_id]),
                            torch.tensor(rollouts.obs[step, :, agent_id]),
                            torch.tensor(
                                rollouts.recurrent_hidden_states[step, :,
                                                                 agent_id]),
                            torch.tensor(
                                rollouts.recurrent_hidden_states_critic[
                                    step, :, agent_id]),
                            torch.tensor(rollouts.masks[step, :, agent_id]),
                            torch.tensor(rollouts.available_actions[step, :,
                                                                    agent_id]))
                    else:
                        actor_critic[agent_id].eval()
                        value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic[
                            agent_id].act(
                                agent_id,
                                torch.tensor(rollouts.share_obs[step, :,
                                                                agent_id]),
                                torch.tensor(rollouts.obs[step, :, agent_id]),
                                torch.tensor(rollouts.recurrent_hidden_states[
                                    step, :, agent_id]),
                                torch.tensor(
                                    rollouts.recurrent_hidden_states_critic[
                                        step, :, agent_id]),
                                torch.tensor(rollouts.masks[step, :,
                                                            agent_id]),
                                torch.tensor(
                                    rollouts.available_actions[step, :,
                                                               agent_id]))

                    values.append(value.detach().cpu().numpy())
                    actions.append(action.detach().cpu().numpy())
                    action_log_probs.append(
                        action_log_prob.detach().cpu().numpy())
                    recurrent_hidden_statess.append(
                        recurrent_hidden_states.detach().cpu().numpy())
                    recurrent_hidden_statess_critic.append(
                        recurrent_hidden_states_critic.detach().cpu().numpy())

            # rearrange action
            actions_env = []
            for i in range(args.n_rollout_threads):
                one_hot_action_env = []
                for agent_id in range(num_agents):
                    one_hot_action = np.zeros(envs.action_space[agent_id].n)
                    one_hot_action[actions[agent_id][i]] = 1
                    one_hot_action_env.append(one_hot_action)
                actions_env.append(one_hot_action_env)

            # Obser reward and next obs
            obs, reward, dones, infos, available_actions = envs.step(
                actions_env)

            # If done then clean the history of observations.
            # insert data in buffer
            masks = []
            for i, done in enumerate(dones):
                mask = []
                for agent_id in range(num_agents):
                    if done:
                        recurrent_hidden_statess[agent_id][i] = np.zeros(
                            args.hidden_size).astype(np.float32)
                        recurrent_hidden_statess_critic[agent_id][
                            i] = np.zeros(args.hidden_size).astype(np.float32)
                        mask.append([0.0])
                    else:
                        mask.append([1.0])
                masks.append(mask)

            bad_masks = []
            high_masks = []
            for info in infos:
                bad_mask = []
                high_mask = []
                for agent_id in range(num_agents):
                    if info[agent_id]['bad_transition']:
                        bad_mask.append([0.0])
                    else:
                        bad_mask.append([1.0])

                    if info[agent_id]['high_masks']:
                        high_mask.append([1.0])
                    else:
                        high_mask.append([0.0])
                bad_masks.append(bad_mask)
                high_masks.append(high_mask)

            if len(envs.observation_space[0]) == 3:
                share_obs = obs.reshape(args.n_rollout_threads, -1,
                                        envs.observation_space[0][1],
                                        envs.observation_space[0][2])
                share_obs = np.expand_dims(share_obs, 1).repeat(num_agents,
                                                                axis=1)

                rollouts.insert(
                    share_obs, obs,
                    np.array(recurrent_hidden_statess).transpose(1, 0, 2),
                    np.array(recurrent_hidden_statess_critic).transpose(
                        1, 0, 2),
                    np.array(actions).transpose(1, 0, 2),
                    np.array(action_log_probs).transpose(1, 0, 2),
                    np.array(values).transpose(1, 0, 2), reward, masks,
                    bad_masks, high_masks, available_actions)
            else:
                share_obs = obs.reshape(args.n_rollout_threads, -1)
                share_obs = np.expand_dims(share_obs, 1).repeat(num_agents,
                                                                axis=1)

                rollouts.insert(
                    share_obs, obs,
                    np.array(recurrent_hidden_statess).transpose(1, 0, 2),
                    np.array(recurrent_hidden_statess_critic).transpose(
                        1, 0, 2),
                    np.array(actions).transpose(1, 0, 2),
                    np.array(action_log_probs).transpose(1, 0, 2),
                    np.array(values).transpose(1, 0, 2), reward, masks,
                    bad_masks, high_masks, available_actions)

        with torch.no_grad():
            for agent_id in range(num_agents):
                if args.share_policy:
                    actor_critic.eval()
                    next_value, _, _ = actor_critic.get_value(
                        agent_id,
                        torch.tensor(rollouts.share_obs[-1, :, agent_id]),
                        torch.tensor(rollouts.obs[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states_critic[-1, :,
                                                                    agent_id]),
                        torch.tensor(rollouts.masks[-1, :, agent_id]))
                    next_value = next_value.detach().cpu().numpy()
                    rollouts.compute_returns(agent_id, next_value,
                                             args.use_gae, args.gamma,
                                             args.gae_lambda,
                                             args.use_proper_time_limits,
                                             args.use_popart,
                                             agents.value_normalizer)
                else:
                    actor_critic[agent_id].eval()
                    next_value, _, _ = actor_critic[agent_id].get_value(
                        agent_id,
                        torch.tensor(rollouts.share_obs[-1, :, agent_id]),
                        torch.tensor(rollouts.obs[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states[-1, :, agent_id]),
                        torch.tensor(
                            rollouts.recurrent_hidden_states_critic[-1, :,
                                                                    agent_id]),
                        torch.tensor(rollouts.masks[-1, :, agent_id]))
                    next_value = next_value.detach().cpu().numpy()
                    rollouts.compute_returns(agent_id, next_value,
                                             args.use_gae, args.gamma,
                                             args.gae_lambda,
                                             args.use_proper_time_limits,
                                             args.use_popart,
                                             agents[agent_id].value_normalizer)

        # update the network
        if args.share_policy:
            actor_critic.train()
            value_loss, action_loss, dist_entropy = agents.update_share(
                num_agents, rollouts)

            logger.add_scalars('reward', {'reward': np.mean(rollouts.rewards)},
                               (episode + 1) * args.episode_length *
                               args.n_rollout_threads)
        else:
            value_losses = []
            action_losses = []
            dist_entropies = []

            for agent_id in range(num_agents):
                actor_critic[agent_id].train()
                value_loss, action_loss, dist_entropy = agents[
                    agent_id].update(agent_id, rollouts)
                value_losses.append(value_loss)
                action_losses.append(action_loss)
                dist_entropies.append(dist_entropy)

                logger.add_scalars(
                    'agent%i/reward' % agent_id,
                    {'reward': np.mean(rollouts.rewards[:, :, agent_id])},
                    (episode + 1) * args.episode_length *
                    args.n_rollout_threads)

        # clean the buffer and reset
        rollouts.after_update()

        total_num_steps = (episode +
                           1) * args.episode_length * args.n_rollout_threads

        if (episode % args.save_interval == 0 or episode == episodes -
                1):  # save for every interval-th episode or for the last epoch
            if args.share_policy:
                torch.save({'model': actor_critic},
                           str(save_dir) + "/agent_model.pt")
            else:
                for agent_id in range(num_agents):
                    torch.save({'model': actor_critic[agent_id]},
                               str(save_dir) + "/agent%i_model" % agent_id +
                               ".pt")

        # log information
        if episode % args.log_interval == 0:
            end = time.time()
            print(
                "\n Map {} Algo {} updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(args.map_name, args.algorithm_name, episode, episodes,
                        total_num_steps, args.num_env_steps,
                        int(total_num_steps / (end - start))))
            if args.share_policy:
                print("value loss of agent: " + str(value_loss))
            else:
                for agent_id in range(num_agents):
                    print("value loss of agent%i: " % agent_id +
                          str(value_losses[agent_id]))

            if args.env_name == "StarCraft2":
                battles_won = []
                battles_game = []
                incre_battles_won = []
                incre_battles_game = []

                for i, info in enumerate(infos):
                    if 'battles_won' in info[0].keys():
                        battles_won.append(info[0]['battles_won'])
                        incre_battles_won.append(info[0]['battles_won'] -
                                                 last_battles_won[i])
                    if 'battles_game' in info[0].keys():
                        battles_game.append(info[0]['battles_game'])
                        incre_battles_game.append(info[0]['battles_game'] -
                                                  last_battles_game[i])

                if np.sum(incre_battles_game) > 0:
                    logger.add_scalars(
                        'incre_win_rate', {
                            'incre_win_rate':
                            np.sum(incre_battles_won) /
                            np.sum(incre_battles_game)
                        }, total_num_steps)
                else:
                    logger.add_scalars('incre_win_rate', {'incre_win_rate': 0},
                                       total_num_steps)
                last_battles_game = battles_game
                last_battles_won = battles_won

        if episode % args.eval_interval == 0 and args.eval:
            eval_battles_won = 0
            eval_episode = 0
            eval_obs, eval_available_actions = eval_env.reset()
            eval_share_obs = eval_obs.reshape(1, -1)
            eval_recurrent_hidden_states = np.zeros(
                (1, num_agents, args.hidden_size)).astype(np.float32)
            eval_recurrent_hidden_states_critic = np.zeros(
                (1, num_agents, args.hidden_size)).astype(np.float32)
            eval_masks = np.ones((1, num_agents, 1)).astype(np.float32)

            while True:
                eval_actions = []
                for agent_id in range(num_agents):
                    if args.share_policy:
                        actor_critic.eval()
                        _, action, _, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic.act(
                            agent_id,
                            torch.tensor(eval_share_obs),
                            torch.tensor(eval_obs[:, agent_id]),
                            torch.tensor(
                                eval_recurrent_hidden_states[:, agent_id]),
                            torch.tensor(
                                eval_recurrent_hidden_states_critic[:,
                                                                    agent_id]),
                            torch.tensor(eval_masks[:, agent_id]),
                            torch.tensor(eval_available_actions[:,
                                                                agent_id, :]),
                            deterministic=True)
                    else:
                        actor_critic[agent_id].eval()
                        _, action, _, recurrent_hidden_states, recurrent_hidden_states_critic = actor_critic[
                            agent_id].act(
                                agent_id,
                                torch.tensor(eval_share_obs),
                                torch.tensor(eval_obs[:, agent_id]),
                                torch.tensor(
                                    eval_recurrent_hidden_states[:, agent_id]),
                                torch.tensor(
                                    eval_recurrent_hidden_states_critic[:,
                                                                        agent_id]
                                ),
                                torch.tensor(eval_masks[:, agent_id]),
                                torch.tensor(
                                    eval_available_actions[:, agent_id, :]),
                                deterministic=True)

                    eval_actions.append(action.detach().cpu().numpy())
                    eval_recurrent_hidden_states[:,
                                                 agent_id] = recurrent_hidden_states.detach(
                                                 ).cpu().numpy()
                    eval_recurrent_hidden_states_critic[:,
                                                        agent_id] = recurrent_hidden_states_critic.detach(
                                                        ).cpu().numpy()

                # rearrange action
                eval_actions_env = []
                for agent_id in range(num_agents):
                    one_hot_action = np.zeros(
                        eval_env.action_space[agent_id].n)
                    one_hot_action[eval_actions[agent_id][0]] = 1
                    eval_actions_env.append(one_hot_action)

                # Obser reward and next obs
                eval_obs, eval_rewards, eval_dones, eval_infos, eval_available_actions = eval_env.step(
                    [eval_actions_env])
                eval_share_obs = eval_obs.reshape(1, -1)

                if eval_dones[0]:
                    eval_episode += 1
                    if eval_infos[0][0]['won']:
                        eval_battles_won += 1
                    for agent_id in range(num_agents):
                        eval_recurrent_hidden_states[0][agent_id] = np.zeros(
                            args.hidden_size).astype(np.float32)
                        eval_recurrent_hidden_states_critic[0][
                            agent_id] = np.zeros(args.hidden_size).astype(
                                np.float32)
                        eval_masks[0][agent_id] = 0.0
                else:
                    for agent_id in range(num_agents):
                        eval_masks[0][agent_id] = 1.0

                if eval_episode >= args.eval_episodes:
                    logger.add_scalars(
                        'eval_win_rate',
                        {'eval_win_rate': eval_battles_won / eval_episode},
                        total_num_steps)
                    break

    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()
    envs.close()
    if args.eval:
        eval_env.close()
コード例 #9
0
def train(env, save_dir: str, epochs: int, update_epochs: int, agents: int,
          trajectory_steps: int, render: bool) -> object:
    reset_dir(save_dir)

    policy = MLPGaussianPolicy(
        env.observation_space.shape[0],
        env.action_space.shape[0],
        hidden_layers=[3, 2],
        action_high=env.action_space.high,
        action_low=env.action_space.low,
    )
    value_fn = MLPValueFn(env.observation_space.shape[0], hidden_layers=[3, 2])
    rl = PPO(policy, value_fn, update_epochs=update_epochs)

    reward_log = DataCSVSaver(
        os.path.join(save_dir, "distance.txt"),
        ("epoch", "averaged reward"))  # log epoch development of reward
    loss_log = DataCSVSaver(
        os.path.join(save_dir, "loss.txt"), ("epoch", "iter", "loss")
    )  # log loss transition to check whether network update is carried out properly

    e = 0
    # training
    while e < epochs:
        try:
            print("Epoch {} ......".format(e))
            rl.reset_sample_buffers()

            average_rewards = []

            # sampling
            print("  sampling...")
            for n in range(agents):
                # buffer to save samples from trajectory
                observations_list = []
                actions_list = []
                rewards_list = []

                # init
                obs = env.reset()
                observations_list.append(obs)

                # run a trajectory
                for t in range(trajectory_steps):
                    action = rl.act(obs)
                    obs, r, done, _ = env.step(action)
                    # save a sample
                    actions_list.append(action)
                    observations_list.append(obs)
                    rewards_list.append(r)
                    if done:
                        break

                rl.feed_trajectory_samples(observations_list, actions_list,
                                           rewards_list, done)

                print(
                    "    agent {}: run for {} steps, average reward {}".format(
                        n, t, np.mean(rewards_list)))
                average_rewards.append(np.mean(rewards_list))

            # update parameters of policy and state value function
            print("  updating...")
            update_epoch_losses = rl.update()

            # logging
            reward_log.append_data(
                e, np.mean(average_rewards))  # save averaged reward
            for i, loss in enumerate(update_epoch_losses):
                loss_log.append_data(e, i,
                                     loss)  # save loss of each update epoch

            print("  average reward {}".format(np.mean(average_rewards)))
            e += 1

        except KeyboardInterrupt:
            command = input("\nSample? Finish? : ")
            if command in ["sample", "Sample"]:
                # run for X steps
                sample_steps = input("How many steps for this sample?: ")
                if sample_steps == "":
                    sample_steps = steps
                    print("default steps {}".format(sample_steps))
                else:
                    sample_steps = int(sample_steps)

                obs = env.reset()
                acc_r = 0
                for t in range(sample_steps):
                    if render:
                        env.render()
                    action = rl.act(obs)
                    obs, r, done, _ = env.step(action)
                    acc_r += r

                continue
            if command in ["finish", "Finish"]:
                print("Ending training ...")
                break

    print("Finish training. Saving the policy and value_fn in {}".format(
        save_dir))
    rl.save(save_dir)
    return rl.policy
コード例 #10
0
def main():
    args = get_config()

    # cuda
    if args.cuda and torch.cuda.is_available():
        device = torch.device("cuda:0")
        torch.set_num_threads(args.n_training_threads)
        if args.cuda_deterministic:
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        device = torch.device("cpu")
        torch.set_num_threads(args.n_training_threads)

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    args.reward_randomization = False

    # path
    model_dir = Path('./results') / args.env_name / args.algorithm_name / (
        "run" + str(args.seed))

    run_dir = model_dir / 'finetune'
    log_dir = run_dir / 'logs'
    save_dir = run_dir / 'models'
    os.makedirs(str(log_dir))
    os.makedirs(str(save_dir))
    logger = SummaryWriter(str(log_dir))

    # env
    envs = make_parallel_env(args)
    #Policy network
    actor_critic = []
    if args.share_policy:
        ac = torch.load(str(model_dir / 'models') +
                        "/agent0_model.pt")['model'].to(device)
        for i in range(args.num_agents):
            actor_critic.append(ac)
    else:
        for i in range(args.num_agents):
            ac = torch.load(
                str(model_dir / 'models') + "/agent%i_model" % i +
                ".pt")['model'].to(device)
            actor_critic.append(ac)
    agents = []
    rollouts = []
    for agent_id in range(args.num_agents):
        # algorithm
        agent = PPO(actor_critic[agent_id],
                    agent_id,
                    args.clip_param,
                    args.ppo_epoch,
                    args.num_mini_batch,
                    args.data_chunk_length,
                    args.value_loss_coef,
                    args.entropy_coef,
                    logger,
                    lr=args.lr,
                    eps=args.eps,
                    max_grad_norm=args.max_grad_norm,
                    use_clipped_value_loss=args.use_clipped_value_loss)

        #replay buffer
        ro = RolloutStorage(args.num_agents, agent_id, args.episode_length,
                            args.n_rollout_threads,
                            envs.observation_space[agent_id],
                            envs.action_space[agent_id],
                            actor_critic[agent_id].recurrent_hidden_state_size)

        agents.append(agent)
        rollouts.append(ro)

    # reset env
    obs = envs.reset()
    # rollout
    for i in range(args.num_agents):
        if len(envs.observation_space[0]) == 1:
            rollouts[i].share_obs[0].copy_(
                torch.tensor(obs.reshape(args.n_rollout_threads, -1)))
            rollouts[i].obs[0].copy_(torch.tensor(obs[:, i, :]))
            rollouts[i].recurrent_hidden_states.zero_()
            rollouts[i].recurrent_hidden_states_critic.zero_()
            rollouts[i].recurrent_c_states.zero_()
            rollouts[i].recurrent_c_states_critic.zero_()
        else:
            raise NotImplementedError
        rollouts[i].to(device)

    # run
    coop_num = []
    defect_num = []
    coopdefect_num = []
    defectcoop_num = []
    gore1_num = []
    gore2_num = []
    collective_return = []
    apple_consumption = []
    waste_cleared = []
    sustainability = []
    fire = []

    start = time.time()
    episodes = int(
        args.num_env_steps) // args.episode_length // args.n_rollout_threads
    all_episode = 0
    turn_on = False

    for episode in range(episodes):
        if (episode > episodes / 2) and (turn_on == False):
            print("\n Turn off fixed actor...")
            # actor
            turn_on = True

        if args.use_linear_lr_decay:
            # decrease learning rate linearly
            for i in range(args.num_agents):
                update_linear_schedule(agents[i].optimizer, episode, episodes,
                                       args.lr)

        for step in range(args.episode_length):
            # Sample actions
            values = []
            actions = []
            action_log_probs = []
            recurrent_hidden_statess = []
            recurrent_hidden_statess_critic = []
            recurrent_c_statess = []
            recurrent_c_statess_critic = []

            with torch.no_grad():
                for i in range(args.num_agents):
                    value, action, action_log_prob, recurrent_hidden_states, recurrent_hidden_states_critic, recurrent_c_states, recurrent_c_states_critic = actor_critic[
                        i].act(
                            rollouts[i].share_obs[step], rollouts[i].obs[step],
                            rollouts[i].recurrent_hidden_states[step],
                            rollouts[i].recurrent_hidden_states_critic[step],
                            rollouts[i].recurrent_c_states[step],
                            rollouts[i].recurrent_c_states_critic[step],
                            rollouts[i].masks[step])
                    values.append(value)
                    actions.append(action)
                    action_log_probs.append(action_log_prob)
                    recurrent_hidden_statess.append(recurrent_hidden_states)
                    recurrent_hidden_statess_critic.append(
                        recurrent_hidden_states_critic)
                    recurrent_c_statess.append(recurrent_c_states)
                    recurrent_c_statess_critic.append(
                        recurrent_c_states_critic)

            # rearrange action
            actions_env = []
            for i in range(args.n_rollout_threads):
                one_hot_action_env = []
                for k in range(args.num_agents):
                    one_hot_action = np.zeros(envs.action_space[0].n)
                    one_hot_action[actions[k][i]] = 1
                    one_hot_action_env.append(one_hot_action)
                actions_env.append(one_hot_action_env)

            # Obser reward and next obs
            obs, reward, done, infos = envs.step(actions_env)

            # If done then clean the history of observations.
            # insert data in buffer
            masks = []
            bad_masks = []
            masks_critic = []
            bad_masks_critic = []
            for i in range(args.num_agents):
                mask = []
                bad_mask = []
                for done_ in done:
                    if done_[i]:
                        mask.append([0.0])
                        bad_mask.append([1.0])
                    else:
                        mask.append([1.0])
                        bad_mask.append([1.0])
                masks.append(torch.FloatTensor(mask))
                bad_masks.append(torch.FloatTensor(bad_mask))

            for i in range(args.num_agents):
                if len(envs.observation_space[0]) == 1:
                    rollouts[i].insert(
                        torch.tensor(obs.reshape(args.n_rollout_threads, -1)),
                        torch.tensor(obs[:,
                                         i, :]), recurrent_hidden_statess[i],
                        recurrent_hidden_statess_critic[i],
                        recurrent_c_statess[i], recurrent_c_statess_critic[i],
                        actions[i], action_log_probs[i], values[i],
                        torch.tensor(reward[:, i].reshape(-1, 1)), masks[i],
                        bad_masks[i])
                else:
                    raise NotImplementedError

        with torch.no_grad():
            next_values = []
            for i in range(args.num_agents):
                next_value = actor_critic[i].get_value(
                    rollouts[i].share_obs[-1], rollouts[i].obs[-1],
                    rollouts[i].recurrent_hidden_states[-1],
                    rollouts[i].recurrent_hidden_states_critic[-1],
                    rollouts[i].recurrent_c_states[-1],
                    rollouts[i].recurrent_c_states_critic[-1],
                    rollouts[i].masks[-1]).detach()
                next_values.append(next_value)

        for i in range(args.num_agents):
            rollouts[i].compute_returns(next_values[i], args.use_gae,
                                        args.gamma, args.gae_lambda,
                                        args.use_proper_time_limits)

        # update the network
        value_losses = []
        action_losses = []
        dist_entropies = []
        for i in range(args.num_agents):
            value_loss, action_loss, dist_entropy = agents[i].update(
                rollouts[i], turn_on)
            value_losses.append(value_loss)
            action_losses.append(action_loss)
            dist_entropies.append(dist_entropy)

        if args.env_name == "StagHunt":
            for info in infos:
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'defect&defect_num' in info.keys():
                    defect_num.append(info['defect&defect_num'])
                if 'coop&defect_num' in info.keys():
                    coopdefect_num.append(info['coop&defect_num'])
                if 'defect&coop_num' in info.keys():
                    defectcoop_num.append(info['defect&coop_num'])

            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'defect&defect_num_per_episode',
                    {'defect&defect_num_per_episode': defect_num[all_episode]},
                    all_episode)
                logger.add_scalars('coop&defect_num_per_episode', {
                    'coop&defect_num_per_episode':
                    coopdefect_num[all_episode]
                }, all_episode)
                logger.add_scalars('defect&coop_num_per_episode', {
                    'defect&coop_num_per_episode':
                    defectcoop_num[all_episode]
                }, all_episode)
                all_episode += 1
        elif args.env_name == "StagHuntGW":
            for info in infos:
                if 'collective_return' in info.keys():
                    collective_return.append(info['collective_return'])
                if 'coop&coop_num' in info.keys():
                    coop_num.append(info['coop&coop_num'])
                if 'gore1_num' in info.keys():
                    gore1_num.append(info['gore1_num'])
                if 'gore2_num' in info.keys():
                    gore2_num.append(info['gore2_num'])

            for i in range(args.n_rollout_threads):
                logger.add_scalars(
                    'collective_return',
                    {'collective_return': collective_return[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'coop&coop_num_per_episode',
                    {'coop&coop_num_per_episode': coop_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore1_num_per_episode',
                    {'gore1_num_per_episode': gore1_num[all_episode]},
                    all_episode)
                logger.add_scalars(
                    'gore2_num_per_episode',
                    {'gore2_num_per_episode': gore2_num[all_episode]},
                    all_episode)

                all_episode += 1

        # clean the buffer and reset
        obs = envs.reset()
        for i in range(args.num_agents):
            if len(envs.observation_space[0]) == 1:
                rollouts[i].share_obs[0].copy_(
                    torch.tensor(obs.reshape(args.n_rollout_threads, -1)))
                rollouts[i].obs[0].copy_(torch.tensor(obs[:, i, :]))
                rollouts[i].recurrent_hidden_states.zero_()
                rollouts[i].recurrent_hidden_states_critic.zero_()
                rollouts[i].recurrent_c_states.zero_()
                rollouts[i].recurrent_c_states_critic.zero_()
            else:
                raise NotImplementedError
            rollouts[i].to(device)

        for i in range(args.num_agents):
            # save for every interval-th episode or for the last epoch
            if (episode % args.save_interval == 0 or episode == episodes - 1):
                torch.save({'model': actor_critic[i]},
                           str(save_dir) + "/agent%i_model" % i + ".pt")

        # log information
        if episode % args.log_interval == 0:
            total_num_steps = (
                episode + 1) * args.episode_length * args.n_rollout_threads
            end = time.time()
            print(
                "\n updates {}/{} episodes, total num timesteps {}/{}, FPS {}.\n"
                .format(episode, episodes, total_num_steps, args.num_env_steps,
                        int(total_num_steps / (end - start))))
            for i in range(args.num_agents):
                print("value loss of agent%i: " % i + str(value_losses[i]))
    logger.export_scalars_to_json(str(log_dir / 'summary.json'))
    logger.close()