def run_experiment(**config): set_seed(config['seed']) original_saved_path = config['saved_path'] if original_saved_path is not None: saved_model = joblib.load(config['saved_path']) if 'config' in saved_model: if not config['override_old_config']: config = saved_model['config'] arguments = { "start_loc": 'all', "include_holdout_obj": False, "persist_goal": config['persist_goal'], "persist_objs": config['persist_objs'], "persist_agent": config['persist_agent'], "feedback_type": config["feedback_type"], "feedback_always": config["feedback_always"], "feedback_freq": config["feedback_freq"], "cartesian_steps": config["cartesian_steps"], "num_meta_tasks": config["rollouts_per_meta_task"], "intermediate_reward": config["intermediate_reward"], } advice_start_index = 160 if original_saved_path is not None: set_seed(config['seed']) policy = saved_model['policy'] optimizer = saved_model['optimizer'] policy.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # TODO: is this necessary? policy.hidden_state = None baseline = saved_model['baseline'] curriculum_step = saved_model['curriculum_step'] env = rl2env(normalize( Curriculum(config['advance_curriculum_func'], start_index=curriculum_step, **arguments)), ceil_reward=config['ceil_reward']) start_itr = saved_model['itr'] reward_predictor = saved_model['reward_predictor'] reward_predictor.hidden_state = None if 'supervised_model' in saved_model: supervised_model = saved_model['supervised_model'] else: supervised_model = None teacher_train_dict = {} for teacher_name in config['feedback_type']: teacher_train_dict[teacher_name] = True else: teacher_train_dict = {} for teacher_name in config['feedback_type']: teacher_train_dict[teacher_name] = True optimizer = None baseline = None env = rl2env(normalize( Curriculum(config['advance_curriculum_func'], start_index=config['level'], **arguments)), ceil_reward=config['ceil_reward']) obs = env.reset() obs_dim = 100 # TODO: consider changing this with 'additional' and adding it! advice_size = sum( [np.prod(obs[adv_k].shape) for adv_k in teacher_train_dict.keys()]) image_dim = 128 memory_dim = config['memory_dim'] instr_dim = config['instr_dim'] use_instr = True instr_arch = 'bigru' use_mem = True arch = 'bow_endpool_res' advice_dim = 128 # TODO: move this to the config policy = ACModel(obs_space=obs_dim, action_space=env.action_space, env=env, image_dim=image_dim, memory_dim=memory_dim, instr_dim=instr_dim, lang_model=instr_arch, use_instr=use_instr, use_memory=use_mem, arch=arch, advice_dim=advice_dim, advice_size=advice_size, num_modules=config['num_modules']) reward_predictor = ACModel( obs_space=obs_dim - 1, # TODO: change into Discrete(3) and do 3-way classification action_space=spaces.Discrete(2), env=env, image_dim=image_dim, memory_dim=memory_dim, instr_dim=instr_dim, lang_model=instr_arch, use_instr=use_instr, use_memory=use_mem, arch=arch, advice_dim=advice_dim, advice_size=advice_size, num_modules=config['num_modules']) if config['self_distill'] and not config['distill_same_model']: obs_dim = env.reset()['obs'].shape[0] image_dim = 128 memory_dim = config['memory_dim'] instr_dim = config['instr_dim'] use_instr = True instr_arch = 'bigru' use_mem = True arch = 'bow_endpool_res' supervised_model = ACModel(obs_space=obs_dim - 1, action_space=env.action_space, env=env, image_dim=image_dim, memory_dim=memory_dim, instr_dim=instr_dim, lang_model=instr_arch, use_instr=use_instr, use_memory=use_mem, arch=arch, advice_dim=advice_dim, advice_size=advice_size, num_modules=config['num_modules']) elif config['self_distill']: supervised_model = policy else: supervised_model = None start_itr = 0 curriculum_step = env.index parser = ArgumentParser() args = parser.parse_args([]) args.entropy_coef = config['entropy_bonus'] args.model = 'default_il' args.lr = config['learning_rate'] args.recurrence = config['backprop_steps'] args.clip_eps = config['clip_eps'] if supervised_model is not None: il_trainer = ImitationLearning( supervised_model, env, args, distill_with_teacher=config['distill_with_teacher']) else: il_trainer = None rp_trainer = ImitationLearning(reward_predictor, env, args, distill_with_teacher=True, reward_predictor=True) teacher_null_dict = env.teacher.null_feedback() obs_preprocessor = make_obs_preprocessor(teacher_null_dict) sampler = MetaSampler( env=env, policy=policy, rollouts_per_meta_task=config['rollouts_per_meta_task'], meta_batch_size=config['meta_batch_size'], max_path_length=config['max_path_length'], parallel=config['parallel'], envs_per_task=1, reward_predictor=reward_predictor, supervised_model=supervised_model, obs_preprocessor=obs_preprocessor, ) sample_processor = RL2SampleProcessor( baseline=baseline, discount=config['discount'], gae_lambda=config['gae_lambda'], normalize_adv=config['normalize_adv'], positive_adv=config['positive_adv'], ) envs = [ copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), copy.deepcopy(env), ] algo = PPOAlgo(policy, envs, config['frames_per_proc'], config['discount'], args.lr, args.beta1, args.beta2, config['gae_lambda'], args.entropy_coef, config['value_loss_coef'], config['max_grad_norm'], args.recurrence, args.optim_eps, config['clip_eps'], config['epochs'], config['meta_batch_size'], parallel=config['parallel'], rollouts_per_meta_task=config['rollouts_per_meta_task'], obs_preprocessor=obs_preprocessor) if optimizer is not None: algo.optimizer.load_state_dict(optimizer) EXP_NAME = get_exp_name(config) exp_dir = os.getcwd() + '/data/' + EXP_NAME + "_" + str(config['seed']) if original_saved_path is None: if os.path.isdir(exp_dir): shutil.rmtree(exp_dir) log_formats = ['stdout', 'log', 'csv'] is_debug = config['prefix'] == 'DEBUG' if not is_debug: log_formats.append('tensorboard') log_formats.append('wandb') logger.configure(dir=exp_dir, format_strs=log_formats, snapshot_mode=config['save_option'], snapshot_gap=50, step=start_itr, name=config['prefix'] + str(config['seed']), config=config) json.dump(config, open(exp_dir + '/params.json', 'w'), indent=2, sort_keys=True, cls=ClassEncoder) advice_end_index, advice_dim = 161, 1 if config[ 'distill_with_teacher']: # TODO: generalize this for multiple feedback types at once! teacher_info = [] else: null_val = np.zeros(advice_end_index - advice_start_index) if len(null_val) > 0: null_val[-1] = 1 teacher_info = [{ "indices": np.arange(advice_start_index, advice_end_index), "null": null_val }] trainer = Trainer( algo=algo, policy=policy, env=deepcopy(env), sampler=sampler, sample_processor=sample_processor, n_itr=config['n_itr'], start_itr=start_itr, success_threshold=config['success_threshold'], accuracy_threshold=config['accuracy_threshold'], exp_name=exp_dir, curriculum_step=curriculum_step, config=config, advance_without_teacher=True, teacher_info=teacher_info, sparse_rewards=not config['intermediate_reward'], distill_only=config['distill_only'], il_trainer=il_trainer, source=config['source'], batch_size=config['meta_batch_size'], train_with_teacher=config['feedback_type'] is not None, distill_with_teacher=config['distill_with_teacher'], supervised_model=supervised_model, reward_predictor=reward_predictor, rp_trainer=rp_trainer, advance_levels=config['advance_levels'], is_debug=is_debug, teacher_train_dict=teacher_train_dict, obs_preprocessor=obs_preprocessor, ) trainer.train()
mean_steps = None for phase_no in range(0, args.phases): logger.info("Starting phase {} with {} demos".format(phase_no, len(train_demos))) if not args.finetune: # Create a new model to be trained from scratch logging.info("Creating new model to be trained from scratch") args.model = model_name + ('_phase_%d' % phase_no) il_learn = ImitationLearning(args) # Train the imitation learning agent if len(train_demos) > 0: il_learn.train(train_demos, writer, csv_writer, status_path, header, reset_status=True) # Stopping criterion valid_log = il_learn.validate(args.val_episodes) success_rate = np.mean([1 if r > 0 else 0 for r in valid_log[0]['return_per_episode']]) if success_rate >= 0.99: logger.info("Reached target success rate with {} demos, stopping".format(len(train_demos))) break eval_seed = grow_training_set(il_learn, train_demos, eval_seed, args.demo_grow_factor, args.num_eval_demos, args.dagger, mean_steps) if __name__ == "__main__": args = parser.parse_args() main(args)
def main(exp, argv): os.environ["BABYAI_STORAGE"] = exp.results_directory() # Parse arguments parser = ArgumentParser() parser.add_argument("--algo", default='ppo', help="algorithm to use (default: ppo)") parser.add_argument("--discount", type=float, default=0.99, help="discount factor (default: 0.99)") parser.add_argument("--reward-scale", type=float, default=20., help="Reward scale multiplier") parser.add_argument( "--gae-lambda", type=float, default=0.99, help="lambda coefficient in GAE formula (default: 0.99, 1 means no gae)" ) parser.add_argument("--value-loss-coef", type=float, default=0.5, help="value loss term coefficient (default: 0.5)") parser.add_argument("--max-grad-norm", type=float, default=0.5, help="maximum norm of gradient (default: 0.5)") parser.add_argument("--clip-eps", type=float, default=0.2, help="clipping epsilon for PPO (default: 0.2)") parser.add_argument("--ppo-epochs", type=int, default=4, help="number of epochs for PPO (default: 4)") parser.add_argument( "--save-interval", type=int, default=50, help= "number of updates between two saves (default: 50, 0 means no saving)") parser.add_argument("--workers", type=int, default=8, help="number of workers for PyTorch (default: 8)") parser.add_argument("--max-count", type=int, default=1000, help="maximum number of frames to run for") parser.add_argument("--sample_duration", type=float, default=0.5, help="sampling duration") parser.add_argument("--cuda", action="store_true", default=False, help="whether to use cuda") args = parser.parse_args(argv) utils.seed(args.seed) torch_settings = init_torch( seed=args.seed, cuda=args.cuda, workers=args.workers, ) # Generate environments envs = [] for i in range(args.procs): env = gym.make(args.env) env.seed(100 * args.seed + i) envs.append(env) # Define model name suffix = datetime.datetime.now().strftime("%y-%m-%d-%H-%M-%S") instr = args.instr_arch if args.instr_arch else "noinstr" mem = "mem" if not args.no_mem else "nomem" model_name_parts = { 'env': args.env, 'algo': args.algo, 'arch': args.arch, 'instr': instr, 'mem': mem, 'seed': args.seed, 'info': '', 'coef': '', 'suffix': suffix } default_model_name = "{env}_{algo}_{arch}_{instr}_{mem}_seed{seed}{info}{coef}_{suffix}".format( **model_name_parts) if args.pretrained_model: default_model_name = args.pretrained_model + '_pretrained_' + default_model_name args.model = args.model.format( **model_name_parts) if args.model else default_model_name utils.configure_logging(args.model) logger = logging.getLogger(__name__) # Define obss preprocessor if 'emb' in args.arch: obss_preprocessor = utils.IntObssPreprocessor( args.model, envs[0].observation_space, args.pretrained_model) else: obss_preprocessor = utils.ObssPreprocessor(args.model, envs[0].observation_space, args.pretrained_model) # Define actor-critic model # acmodel = utils.load_model(args.model, raise_not_found=False) acmodel = None if acmodel is None: if args.pretrained_model: acmodel = utils.load_model(args.pretrained_model, raise_not_found=True) else: acmodel = ACModel(obss_preprocessor.obs_space, envs[0].action_space, args.image_dim, args.memory_dim, args.instr_dim, not args.no_instr, args.instr_arch, not args.no_mem, args.arch) obss_preprocessor.vocab.save() # utils.save_model(acmodel, args.model) if torch_settings.cuda: acmodel.cuda() # Define actor-critic algo reshape_reward = lambda _0, _1, reward, _2: args.reward_scale * reward if args.algo == "ppo": algo = babyai.rl.PPOAlgo( envs, acmodel, args.frames_per_proc, args.discount, args.lr, args.beta1, args.beta2, args.gae_lambda, args.entropy_coef, args.value_loss_coef, args.max_grad_norm, args.recurrence, args.optim_eps, args.clip_eps, args.ppo_epochs, args.batch_size, obss_preprocessor, reshape_reward) else: raise ValueError("Incorrect algorithm name: {}".format(args.algo)) # When using extra binary information, more tensors (model params) are initialized compared to when we don't use that. # Thus, there starts to be a difference in the random state. If we want to avoid it, in order to make sure that # the results of supervised-loss-coef=0. and extra-binary-info=0 match, we need to reseed here. utils.seed(args.seed) # Restore training status status_path = os.path.join(utils.get_log_dir(args.model), 'status.json') if os.path.exists(status_path): with open(status_path, 'r') as src: status = json.load(src) else: status = {'i': 0, 'num_episodes': 0, 'num_frames': 0} # # Define logger and Tensorboard writer and CSV writer # header = (["update", "episodes", "frames", "FPS", "duration"] # + ["return_" + stat for stat in ['mean', 'std', 'min', 'max']] # + ["success_rate"] # + ["num_frames_" + stat for stat in ['mean', 'std', 'min', 'max']] # + ["entropy", "value", "policy_loss", "value_loss", "loss", "grad_norm"]) # if args.tb: # from tensorboardX import SummaryWriter # writer = SummaryWriter(utils.get_log_dir(args.model)) # csv_path = os.path.join(utils.get_log_dir(args.model), 'log.csv') # first_created = not os.path.exists(csv_path) # # we don't buffer data going in the csv log, cause we assume # # that one update will take much longer that one write to the log # csv_writer = csv.writer(open(csv_path, 'a', 1)) # if first_created: # csv_writer.writerow(header) # Log code state, command, availability of CUDA and model babyai_code = list(babyai.__path__)[0] try: last_commit = subprocess.check_output( 'cd {}; git log -n1'.format(babyai_code), shell=True).decode('utf-8') logger.info('LAST COMMIT INFO:') logger.info(last_commit) except subprocess.CalledProcessError: logger.info('Could not figure out the last commit') try: diff = subprocess.check_output('cd {}; git diff'.format(babyai_code), shell=True).decode('utf-8') if diff: logger.info('GIT DIFF:') logger.info(diff) except subprocess.CalledProcessError: logger.info('Could not figure out the last commit') logger.info('COMMAND LINE ARGS:') logger.info(args) logger.info("CUDA available: {}".format(torch.cuda.is_available())) logger.info(acmodel) # Train model total_start_time = time.time() best_success_rate = 0 best_mean_return = 0 test_env_name = args.env wrapper = iteration_wrapper( exp, sync=torch_settings.sync, max_count=args.max_count, sample_duration=args.sample_duration, ) # while status['num_frames'] < args.frames: while True: with wrapper() as it: # Update parameters if wrapper.done(): break update_start_time = time.time() logs = algo.update_parameters() update_end_time = time.time() it.set_count(logs["num_frames"]) it.log(loss=logs["loss"], )