How to switch between policies in multiagent rllib? - tensorflow

Question

For my thesis, I want to train a PPO model on a collaborative multi-agent situation. The 'teammates' are already trained Behavior Cloning models each representing different kind of behavior. These models will not be trained simultaneously to the PPO. I want the PPO encounter all these different models during training, meaning: I want to switch out the 'teammate' of the PPO every x amount of epochs. How do I do this?
I am using the overcooked-ai environment. All code would be too much to include. But the generation of the trainer ends up returning this:
trainer = PPOTrainer(env="overcooked_multi_agent", config={
"multiagent": multi_agent_config,
"callbacks" : TrainingCallbacks,
"custom_eval_function" : get_rllib_eval_function(evaluation_params, environment_params['eval_mdp_params'], environment_params['env_params'],
environment_params["outer_shape"], 'ppo', 'ppo' if self_play else 'bc',
verbose=params['verbose']),
"env_config" : environment_params,
"eager" : False,
**training_params
}, logger_creator=custom_logger_creator
and the multi-agent-config is created like
def gen_policy(policy_type="ppo"):
# supported policy types thus far
assert policy_type in ["ppo", "bc"]
if policy_type == "ppo":
config = {
"model" : {
"custom_options" : model_params,
"custom_model" : "MyPPOModel"
}
}
return (None, env.ppo_observation_space, env.action_space, config)
elif policy_type == "bc":
bc_cls = bc_params['bc_policy_cls']
bc_config = bc_params['bc_config']
return (bc_cls, env.bc_observation_space, env.action_space, bc_config)
# Rllib compatible way of setting the directory we store agent checkpoints in
logdir_prefix = "{0}_{1}_{2}".format(params["experiment_name"], params['training_params']['seed'], timestr)
def custom_logger_creator(config):
"""Creates a Unified logger that stores results in <params['results_dir']>/<params["experiment_name"]>_<seed>_<timestamp>
"""
results_dir = params['results_dir']
if not os.path.exists(results_dir):
try:
os.makedirs(results_dir)
except Exception as e:
print("error creating custom logging dir. Falling back to default logdir {}".format(DEFAULT_RESULTS_DIR))
results_dir = DEFAULT_RESULTS_DIR
logdir = tempfile.mkdtemp(
prefix=logdir_prefix, dir=results_dir)
logger = UnifiedLogger(config, logdir, loggers=None)
return logger
# Create rllib compatible multi-agent config based on params
multi_agent_config = {}
all_policies = ['ppo']
# Whether both agents should be learned
self_play = iterable_equal(multi_agent_params['bc_schedule'], OvercookedMultiAgent.self_play_bc_schedule)
if not self_play:
all_policies.append('bc')
multi_agent_config['policies'] = { policy : gen_policy(policy) for policy in all_policies }
def select_policy(agent_id):
if agent_id.startswith('ppo'):
return 'ppo'
if agent_id.startswith('bc'):
return 'bc'
multi_agent_config['policy_mapping_fn'] = select_policy
multi_agent_config['policies_to_train'] = 'ppo'
Do I already change my configuration to include multiple BHC models (code above is built to handle 1), or can I somehow change the models in the trainingloop?
Any ideas would be super helpful!