general_exploit.py

from mars.utils.common import SelfplayBasedMethods
from mars.env.import_env import make_env
from mars.rollout import rollout
from mars.rl.agents  import *
from mars.rl.agents.multiagent import MultiAgent
from mars.utils.func import get_model_path, get_exploiter
from mars.utils.args_parser import get_args

def launch():
    args = get_args()
    env = '_'.join([args.env_type, args.env_name])
    method = args.marl_method

    # exploitation setting
    args.num_envs = 1
    args.max_episodes = 10000  # exploitation episodes
    args.against_baseline = False
    args.test = False
    # args.render = True
    args.exploit = True

    folder = f'./data/model/{args.load_id}/{env}_{method}/'
    args.load_model_full_path = get_model_path(method, folder)

    ### Create env
    env = make_env(args)
    print(env)

    ### Specify models for each agent     
    trained_model = eval(args.algorithm)(env, args)
    
    ### Load exploiter with specified args (just change the previous args)
    exploiter, exploitation_args = get_exploiter('PPO', env, args)  # PPO can handle continuous action 

    ### Construct multi-agent model
    if 'first' in args.to_exploit:
        exploitation_args['idx_exploited_model'] = 0  # exploit the first player side
        model = MultiAgent(env, [trained_model, exploiter], exploitation_args)

    else: # 'second'
        exploitation_args['idx_exploited_model'] = 1  # exploit the second player side
        model = MultiAgent(env, [exploiter, trained_model], exploitation_args)

    ### Rollout
    rollout(env, model, exploitation_args, save_id = str(args.load_id)+f'_exploit_{args.to_exploit}') # save results of exploitation in a separate folder

if __name__ == '__main__':
    launch()