def train(config, dir_manager=None, logger=None, pbar="default_pbar"): # A few safety checks check_training_args(config) # Creates a directory manager that encapsulates our directory-tree structure if dir_manager is None: dir_manager = DirectoryManager(agent_alg=config.agent_alg, env_name=config.env_name, desc=config.desc, seed=config.seed) dir_manager.create_directories() # Creates logger and prints config if logger is None: logger = create_logger('MASTER', config.log_level, dir_manager.seed_dir / 'logger.out') logger.debug(config_to_str(config)) # Creates a progress-bar if type(pbar) is str: if pbar == "default_pbar": pbar = tqdm() if pbar is not None: pbar.n = 0 pbar.desc += f'{dir_manager.storage_dir.name}/{dir_manager.experiment_dir.name}/{dir_manager.seed_dir.name}' pbar.total = config.n_episodes # Encapsulates in a dict all user-defined params that concern the world (scenario.make_world()) world_params = {} world_params['use_dense_rewards'] = config.use_dense_rewards if config.env_name == 'chase': if config.n_preys is not None: world_params['n_preys'] = config.n_preys if config.n_preds is not None: world_params['n_preds'] = config.n_preds if config.prey_variance is not None: world_params['prey_variance'] = config.prey_variance if config.individual_reward is not None: world_params['individual_reward'] = config.individual_reward elif config.env_name == 'gather': if config.n_agents is not None: world_params['n_agents'] = config.n_agents elif config.env_name == 'intersection': if config.n_agents is not None: world_params['n_agents'] = config.n_agents elif config.env_name == 'bounce': world_params['episode_length'] = config.episode_length if config.line_length is not None: world_params['line_length'] = config.line_length elif config.env_name == 'compromise': if config.line_length is not None: world_params['line_length'] = config.line_length if config.show_all_landmarks is not None: world_params['show_all_landmarks'] = config.show_all_landmarks elif config.env_name == 'imitation': if config.staged is not None: world_params['staged'] = config.staged if config.set_trap is not None: world_params['set_trap'] = config.set_trap elif config.env_name == 'intersection': if config.by_stander is not None: world_params['by_stander'] = config.by_stander elif config.env_name == 'spread': if config.n_agents is not None: world_params['n_agents'] = config.n_agents if config.shuffle_landmarks is not None: world_params['shuffle_landmarks'] = config.shuffle_landmarks if config.color_objects is not None: world_params['color_objects'] = config.color_objects if config.small_agents is not None: world_params['small_agents'] = config.small_agents save_dict_to_json(world_params, str(dir_manager.seed_dir / 'world_params.json')) # Encapsulates in a dict all user-defined params that concern the environment (multiagent.environment.MultiAgentEnv) env_params = {} env_params['env_name'] = config.env_name if 'football' not in config.env_name: env_params['use_max_speed'] = config.use_max_speed save_dict_to_json(env_params, str(dir_manager.seed_dir / 'env_params.json')) # Sets the random seeds (for reproducibility) set_seeds(config.seed) # Initializes environments # TODO: Check reproductibility and that different envs are seeded differently if '3v2football' == config.env_name: obs_rep = config.representation if config.feature_extractor == 'identity': assert obs_rep in ['simple115', 'simple37'] elif config.feature_extractor == 'convNet': assert obs_rep == 'extracted' else: raise NotImplemented( f"config.feature_extractor={config.feature_extractor} not recognized." ) env = make_parallel_football_env( seed_dir=dir_manager.seed_dir, seed=config.seed, dump_freq=config.dump_freq, representation=obs_rep, render=False, n_rollout_threads=config.n_rollout_threads ) # no rendering during training else: env = make_parallel_particle_env( scenario_name=config.env_name, n_rollout_threads=config.n_rollout_threads, seed=config.seed, use_discrete_action=config.use_discrete_action, use_max_speed=config.use_max_speed, world_params=world_params) if not config.use_cuda: torch.set_num_threads(config.n_training_threads) # Initialize the algo algorithm = init_from_config(env, config, logger) # Creates recorders and stores basic info regarding agent types os.makedirs(dir_manager.recorders_dir, exist_ok=True) train_recorder = algorithm.create_train_recorder() train_recorder.tape['agent_colors'] = env.agent_colors if 'football' in config.env_name: if config.feature_extractor == "convNet": n_stack = 4 elif config.feature_extractor == "identity": n_stack = 1 else: raise NotImplemented obs_buffers = ObsBufferCollection(n_env=config.n_rollout_threads, n_stack=n_stack) replay_buffer = StackingReplayBuffer( max_steps=config.buffer_length, num_agents=algorithm.nagents, obs_dims=[obsp.shape for obsp in env.observation_space], ac_dims=[ acsp.shape[0] if isinstance(acsp, Box) else acsp.n for acsp in env.action_space ], n_stack=n_stack) else: # defines observation buffer for multi-step obs_buffers = ObsBufferCollection(n_env=config.n_rollout_threads, n_stack=1) replay_buffer = ReplayBuffer( max_steps=config.buffer_length, num_agents=algorithm.nagents, obs_dims=[obsp.shape for obsp in env.observation_space], ac_dims=[ acsp.shape[0] if isinstance(acsp, Box) else acsp.n for acsp in env.action_space ]) # Saves initial models current_model = "model_ep0.pt" best_eval_reward_exploit = -100000. best_model_exploit = "model_ep0_exploit_best.pt" algorithm.save(dir_manager.seed_dir / current_model) algorithm.save(dir_manager.seed_dir / best_model_exploit) best_eval_reward_explore = -100000. best_model_explore = "model_ep0_explore_best.pt" algorithm.save(dir_manager.seed_dir / current_model) algorithm.save(dir_manager.seed_dir / best_model_explore) # Initializes step and episode counters step_i = 0 ep_steps = np.zeros(shape=(config.n_rollout_threads, ), dtype=np.int) ep_dones = 0 ep_recorders = [ EpisodeRecorder(stuff_to_record=['reward']) for _ in range(config.n_rollout_threads) ] obs = env.reset() obs_buffers.fill(obs) algorithm.set_exploration( begin_decay_proportion=config.begin_exploration_decay, n_episodes=config.n_episodes, end_decay_proportion=config.end_exploration_decay, initial_scale=config.init_noise_scale, final_scale=config.final_noise_scale, current_episode=ep_dones) # EPISODES LOOP while ep_dones < config.n_episodes: start_time = time.time() # ENVIRONMENT STEP # convert observations to torch Variable torch_obs = [ Variable(torch.Tensor(obs_buffers.read()[:, i]), requires_grad=False) for i in range(algorithm.nagents) ] # get actions as torch Variables torch_agent_actions = algorithm.select_action(torch_obs, is_exploring=True) # convert actions to numpy arrays agent_actions = [ac.data.numpy() for ac in torch_agent_actions] # rearrange actions to be per environment actions = [[ac[i] for ac in agent_actions] for i in range(config.n_rollout_threads)] # makes one step in the environment next_obs, rewards, dones, infos = env.step(actions) # put transitions in the memory buffer replay_buffer.push(obs, agent_actions, rewards, next_obs, dones) # saves relevant info in episode recorders for i in range(config.n_rollout_threads): ep_recorders[i].add_step(obs[i], actions[i], rewards[i], next_obs[i]) # ending step obs = next_obs obs_buffers.append(obs) step_i += config.n_rollout_threads step_time = time.time() - start_time ep_steps += 1 # LEARNING STEP if (len(replay_buffer) >= config.batch_size * config.warmup) \ and (step_i % config.steps_per_update) < config.n_rollout_threads: # Prepares models to training if config.use_cuda: algorithm.prep_training(device='gpu') else: algorithm.prep_training(device='cpu') # Performs one algorithm update sample = replay_buffer.sample(config.batch_size, to_gpu=config.use_cuda, normalize_rewards=False) algorithm.update(sample, train_recorder) # Update target networks algorithm.update_all_targets() # Prepares models to go back in rollout phase algorithm.prep_rollouts(device='cpu') # EPISODE ENDINGS episodes_over = dones | (ep_steps >= config.episode_length) if any(episodes_over): if pbar is not None: pbar.update(sum(episodes_over)) for env_i, is_over in enumerate(episodes_over): if is_over: ep_dones += 1 ep_steps[env_i] = 0 # Reset environments obs[env_i] = env.reset(env_i=env_i) obs_buffers[env_i].flush() obs_buffers[env_i].fill(obs[env_i]) # Summarizes episode metrics train_recorder.append( 'total_reward', ep_recorders[env_i].get_total_reward()) # Reinitialise episode recorder ep_recorders[env_i] = EpisodeRecorder( stuff_to_record=['reward']) # Printing if one third of training is completed if (ep_dones - 1) % (config.n_episodes // 3) == 0 and ep_dones != config.n_episodes: step_time = time.time() - start_time logger.info( f"Episode {ep_dones}/{config.n_episodes}, " f"speed={round_to_two(float(config.n_rollout_threads) / step_time)}steps/s" ) # Sets exploration noise current_noise_scale = algorithm.set_exploration( begin_decay_proportion=config.begin_exploration_decay, n_episodes=config.n_episodes, end_decay_proportion=config.end_exploration_decay, initial_scale=config.init_noise_scale, final_scale=config.final_noise_scale, current_episode=ep_dones) # BOOK-KEEPING if ep_dones % config.episodes_per_save < config.n_rollout_threads: # Model checkpoints if config.save_incrementals: os.makedirs(dir_manager.incrementals_dir, exist_ok=True) algorithm.save(dir_manager.incrementals_dir / ('model_ep%i.pt' % (ep_dones + 1))) os.remove(dir_manager.seed_dir / current_model) current_model = f"model_ep{ep_dones}.pt" algorithm.save(dir_manager.seed_dir / current_model) logger.debug('Saving model checkpoint') # Current model evaluation (run episodes without exploration) if config.n_evaluation_episodes > 0: logger.debug( f'Evaluating model for {config.n_evaluation_episodes} episodes' ) set_seeds( config.evaluation_seed) # fixed seed for evaluation env.seed(config.evaluation_seed) eval_config = get_evaluation_args(overwritten_args="") eval_config.storage_name = dir_manager.storage_dir.name eval_config.experiment_num = int( dir_manager.experiment_dir.stem.strip('experiment')) eval_config.seed_num = int( dir_manager.seed_dir.stem.strip('seed')) eval_config.render = False eval_config.n_episodes = config.n_evaluation_episodes eval_config.last_model = True eval_config.noise_scale = None eval_config.episode_length = config.episode_length eval_config.representation = config.representation # Evaluate with exploit (without explorarion) eval_reward_exploit = np.vstack(evaluate(eval_config)) train_recorder.append('eval_episodes', ep_dones) train_recorder.append('eval_total_reward_exploit', eval_reward_exploit) if eval_reward_exploit.mean() > best_eval_reward_exploit: logger.debug("New best exploit model") os.remove(dir_manager.seed_dir / best_model_exploit) best_model_exploit = f"model_ep{ep_dones}_exploit_best.pt" algorithm.save(dir_manager.seed_dir / best_model_exploit) best_eval_reward_exploit = eval_reward_exploit.mean() # Evaluate with exploration eval_config.noise_scale = current_noise_scale eval_reward_explore = np.vstack(evaluate(eval_config)) train_recorder.append('eval_total_reward_explore', eval_reward_explore) if eval_reward_explore.mean() > best_eval_reward_explore: logger.debug("New best explore model") os.remove(dir_manager.seed_dir / best_model_explore) best_model_explore = f"model_ep{ep_dones}_explore_best.pt" algorithm.save(dir_manager.seed_dir / best_model_explore) best_eval_reward_explore = eval_reward_explore.mean() set_seeds(config.seed + ep_dones) env.seed(config.seed + ep_dones) # Graphs checkpoints logger.debug('Saving recorder checkpoints and graphs') train_recorder.save(dir_manager.recorders_dir / 'train_recorder.pkl') # Saving graphs if len(train_recorder.tape['actor_loss']) > 0: algorithm.save_training_graphs( train_recorder=train_recorder, save_dir=dir_manager.seed_dir) # Saves model one last time and close the environment os.remove(dir_manager.seed_dir / current_model) current_model = f"model_ep{ep_dones}.pt" algorithm.save(dir_manager.seed_dir / current_model) env.close()
def _compute_seed_scores(storage_dir, performance_metric, performance_aggregation, group_key, bar_key, re_run_if_exists, save_dir, logger, root_dir, n_eval_runs): if (storage_dir / save_dir / f"{save_dir}_seed_scores.pkl").exists() and not re_run_if_exists: logger.info( f" SKIPPING {storage_dir} - {save_dir}_seed_scores.pkl already exists" ) return else: logger.info(f"Benchmarking {storage_dir}...") assert group_key in [ 'task_name', 'storage_name', 'experiment_num', 'alg_name' ] assert bar_key in [ 'task_name', 'storage_name', 'experiment_num', 'alg_name' ] # Initialize container scores = OrderedDict() # Get all experiment directories all_experiments = DirectoryTree.get_all_experiments( storage_dir=storage_dir) for experiment_dir in all_experiments: # For that experiment, get all seed directories experiment_seeds = DirectoryTree.get_all_seeds( experiment_dir=experiment_dir) # Initialize container all_seeds_scores = [] for i, seed_dir in enumerate(experiment_seeds): # Prints which seed directory is being treated logger.debug(f"{seed_dir}") # Loads training config config_dict = load_dict_from_json(str(seed_dir / "config.json")) # Selects how data will be identified keys = { "task_name": config_dict["task_name"], "storage_name": seed_dir.parents[1].name, "alg_name": config_dict["alg_name"], "experiment_num": seed_dir.parents[0].name.strip('experiment') } outer_key = keys[bar_key] inner_key = keys[group_key] # Evaluation phase if performance_metric == 'evaluation_runs': assert n_eval_runs is not None try: from evaluate import evaluate, get_evaluation_args except ImportError as e: raise ImportError( f"{e}\nTo evaluate models based on --performance_metric='evaluation_runs' " f"alfred.benchmark assumes the following structure that the working directory contains " f"a file called evaluate.py containing two functions: " f"\n\t1. a function evaluate() that returns a score for each evaluation run" f"\n\t2. a function get_evaluation_args() that returns a Namespace of arguments for evaluate()" ) # Sets config for evaluation phase eval_config = get_evaluation_args(overwritten_args="") eval_config.storage_name = seed_dir.parents[1].name eval_config.experiment_num = int( seed_dir.parents[0].name.strip("experiment")) eval_config.seed_num = int(seed_dir.name.strip("seed")) eval_config.render = False eval_config.n_episodes = n_eval_runs eval_config.root_dir = root_dir # Evaluates agent and stores the return performance_data = evaluate(eval_config) else: # Loads training data loaded_recorder = Recorder.init_from_pickle_file( filename=str(seed_dir / 'recorders' / 'train_recorder.pkl')) performance_data = loaded_recorder.tape[performance_metric] # Aggregation phase if performance_aggregation == 'min': score = np.min(performance_data) elif performance_aggregation == 'max': score = np.max(performance_data) elif performance_aggregation == 'avg': score = np.mean(performance_data) elif performance_aggregation == 'last': score = performance_data[-1] elif performance_aggregation == 'mean_on_last_20_percents': eighty_percent_index = int(0.8 * len(performance_data)) score = np.mean(performance_data[eighty_percent_index:]) else: raise NotImplementedError all_seeds_scores.append(score) if outer_key not in scores.keys(): scores[outer_key] = OrderedDict() scores[outer_key][inner_key] = np.stack(all_seeds_scores) os.makedirs(storage_dir / save_dir, exist_ok=True) with open(storage_dir / save_dir / f"{save_dir}_seed_scores.pkl", "wb") as f: pickle.dump(scores, f) scores_info = { 'n_eval_runs': n_eval_runs, 'performance_metric': performance_metric, 'performance_aggregation': performance_aggregation } save_dict_to_json(scores_info, filename=str(storage_dir / save_dir / f"{save_dir}_seed_scores_info.json"))