<reponame>maxiaoba/rlk
import torch
import numpy as np
import time
import pdb
from rlkit.torch.core import eval_np
from rlkit.torch.policies.tanh_gaussian_policy import TanhGaussianPolicy
from rlkit.torch.policies.make_deterministic import MakeDeterministic
from rlkit.torch.policies.gumbel_softmax_policy import GumbelSoftmaxMlpPolicy
from rlkit.policies.argmax import ArgmaxDiscretePolicy

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--exp_name', type=str, default='simple_spread')
parser.add_argument('--boundary', action='store_true', default=False)
parser.add_argument('--sb', action='store_true', default=False) # shared obs
parser.add_argument('--num_ag', type=int, default=None)
parser.add_argument('--num_adv', type=int, default=None)
parser.add_argument('--num_l', type=int, default=None)
parser.add_argument('--ar', type=float, default=None) # agent radius
parser.add_argument('--sr', action='store_true', default=False) # shape reward
parser.add_argument('--sg', action='store_true', default=False) # shared groups
parser.add_argument('--mpl', type=int, default=25) # max path length
parser.add_argument('--log_dir', type=str, default='MASAC')
parser.add_argument('--epoch', type=int, default=None)
parser.add_argument('--seed', type=str, default=0)
args = parser.parse_args()

pre_dir = './Data/'+args.exp_name\
            +('bd' if args.boundary else '')\
            +('sr' if args.sr else '')\
            +(('ag'+str(args.num_ag)) if args.num_ag else '')\
            +(('adv'+str(args.num_adv)) if args.num_adv else '')\
            +(('l'+str(args.num_l)) if args.num_l else '')\
            +(('ar'+str(args.ar)) if args.ar else '')\
            +'_mpl'+str(args.mpl)
if args.epoch:
    data_path = '{}/{}/seed{}/itr_{}.pkl'.format(pre_dir,args.log_dir,args.seed,args.epoch)
else:
    data_path = '{}/{}/seed{}/params.pkl'.format(pre_dir,args.log_dir,args.seed)
data = torch.load(data_path,map_location='cpu')
policy_n = data['trainer/policy_n']
if isinstance(policy_n[0],TanhGaussianPolicy):
    policy_n = [MakeDeterministic(policy) for policy in policy_n]
elif  isinstance(policy_n[0],GumbelSoftmaxMlpPolicy):
    policy_n = [ArgmaxDiscretePolicy(policy,use_preactivation=True) for policy in policy_n]

if 'trainer/shared_gnn' in data.keys():
    shared_gnn = data['trainer/shared_gnn']
elif 'R2GGNN12Share' in args.log_dir:
    shared_gnn = data['trainer/qf1'].obs_gnn
else:
    shared_gnn = None

world_args=dict(
    num_agents=args.num_ag,
    num_adversaries=args.num_adv,
    num_landmarks=args.num_l,
    agent_size=args.ar,
    boundary=([[-1.,-1.],[1.,1.]] if args.boundary else None),
    shape_reward=args.sr,
    shared_obs=args.sb,
)
import sys
sys.path.append("./particle-graph-envs")
from make_env import make_env
env = make_env(args.exp_name,discrete_action_space=False,world_args=world_args)
num_agent = env.num_agents
obs_dim = env.observation_space.low.size
action_dim = env.action_space.low.size

if args.sg:
    from build_networks import build_shared_groups
    shared_groups = build_shared_groups(args)
else:
    shared_groups = dict()
    for i in range(num_agent):
        shared_groups[i] = 0
print('shared_groups: ', shared_groups)

max_path_length = args.mpl
path_length = 0
done = np.array([False]*num_agent)
c_r = np.zeros(num_agent)
o_n = env.reset()
while True:
    path_length += 1
    a_n = []
    if args.sb:
        if shared_gnn:
            if args.sg:
                o_emb_n = [eval_np(sub_shared_encoder,o_n[None,:])[0] for sub_shared_encoder in shared_gnn]
                for i, policy in enumerate(policy_n):
                    o = o_emb_n[shared_groups[i]][i]
                    a, info = policy.get_action(o)
                    a_n.append(a)
            else:
                if isinstance(shared_gnn, list):
                    assert len(shared_gnn) == 1
                    shared_gnn = shared_gnn[0]
                o_emb_n = eval_np(shared_gnn,o_n[None,:])[0]
                for i,policy in enumerate(policy_n):
                    o = o_emb_n[i]
                    a, info = policy.get_action(o)
                    a_n.append(a)
        else:
            for policy in policy_n:
                a, info = policy.get_action(o_n)
                a_n.append(a)
    else:
        for o,policy in zip(o_n,policy_n):
            a, info = policy.get_action(o)
            a_n.append(a)

    o_n, r_n, done, _ = env.step(a_n)
    c_r += r_n
    env.render()
    print("step: ",path_length)
    print("a: ",a_n)
    print("info: ",info)
    print("o: ",o_n)
    print('r: ',r_n)
    print(done)
    time.sleep(0.1)
    if path_length > max_path_length or done.all():
        print('c_r: ',c_r)
        path_length = 0
        done = np.array([False]*num_agent)
        c_r = np.zeros(num_agent)
        o_n = env.reset()
        pdb.set_trace()