import ray.rllib.agents.ppo as ppo
import ray.rllib.agents.a3c as a3c
import ray.rllib.agents.impala as impala
# import ..agents.ppo as ppo
# import ..agents.a3c as a3c
# import ..agents.impala as impala


def init_trainer(name, config):
    if name == 'a3c':
        trainer = a3c.A3CTrainer(config=config)
    elif name == 'impala':
        trainer = impala.ImpalaTrainer(config=config)
    elif name == 'ppo':
        trainer = ppo.PPOTrainer(config=config)
    return trainer


def get_default_config(name):
    if name == 'a3c':
        config = a3c.DEFAULT_CONFIG.copy()
    elif name == 'impala':
        config = impala.DEFAULT_CONFIG.copy()
    elif name == 'ppo':
        config = ppo.DEFAULT_CONFIG.copy()
    return config
