def train_sureal(): P.initialize_experiment() ctx = mp.get_context("spawn") pipe_rl_end, pipe_sup_end = ctx.Pipe() rlsup_params = P.get_current_parameters()["RLSUP"] sim_seed_dataset = rlsup_params.get("sim_seed_dataset") run_name = P.get_current_parameters()["Setup"]["run_name"] # Trigger rebuild if necessary before going into all the threads and processes _ = get_restricted_env_id_lists() _ = get_restricted_env_id_lists(full=True) if sim_seed_dataset: copy_seed_dataset(from_dataset=sim_seed_dataset, to_dataset=rl_dataset_name(run_name)) if DEBUG_SUP: train_supervised_worker(pipe_sup_end) elif DEBUG_RL: # Start supervised learning in another process. Keep RL in main process. sup_proces = ctx.Process(target=train_supervised_worker, args=[pipe_sup_end]) sup_proces.start() train_rl_worker(pipe_rl_end) else: rl_process = ctx.Process(target=train_rl_worker, args=[pipe_rl_end]) sup_proces = ctx.Process(target=train_supervised_worker, args=[pipe_sup_end]) rl_process.start() sup_proces.start()
def generate_rollout_amt_visualizations(): setup = P.get_current_parameters()["Setup"] dataset_name = setup.get("viz_dataset_name") or get_eval_tmp_dataset_name( setup["model"], setup["run_name"]) print(f"Generating AMT animations for dataset: {dataset_name}") pic_domain = "sim" data_domain = "real" # Some quick params. TODO: Bring this into json viz_params = { "draw_drone": True, "draw_trajectory": True, "draw_fov": True, "include_vdist": False, "include_layer": None, "include_instr": False } print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() # TODO: Grab the correct env list env_list = test_envs viz = RolloutVisualizer(resolution=400) base_dir = os.path.join(get_rollout_viz_dir(), f"{dataset_name}-{data_domain}") os.makedirs(base_dir, exist_ok=True) for env_id in env_list: try: env_data = load_single_env_from_dataset(dataset_name, env_id, "supervised") except FileNotFoundError as e: print(f"Skipping env: {env_id}") continue if len(env_data) == 0: print(f"Skipping env: {env_id}. Rollout exists but is EMPTY!") continue segs = split_into_segs(env_data) for seg in segs: seg_idx = seg[0]["seg_idx"] seg_name = f"{env_id}:0:{seg_idx}-{data_domain}" gif_filename = f"{seg_name}-roll.gif" instr_filename = f"{seg_name}-instr.txt" # Generate and save gif frames = viz.top_down_visualization(env_id, seg_idx, seg, pic_domain, viz_params) print("Saving GIF") viz.presenter.save_gif(frames, os.path.join(base_dir, gif_filename), fps=5.0) # Save instruction with open(os.path.join(base_dir, instr_filename), "w") as fp: fp.write(seg[0]["instruction"]) print("ding")
def build_noisy_pose_data(): """ Randomly sample pose noise for every observation in every environment for the RSS experiment with noisy poses. This needs to be pre-computed once before training to simulate the noise being measured during trajectory collection. If we were to randomize poses during training, that would be akin to regularization, which could actually improve instead of hurt performance. :return: """ initialize_experiment() params = get_current_parameters() setup_params = params["Setup"] train_envs, dev_envs, test_envs = get_restricted_env_id_lists( max_envs=setup_params["max_envs"], prune_ambiguous=setup_params["prune_ambiguous"]) envs = dev_envs + train_envs + test_envs print("Num envs:" + str(len(envs))) pos_noise = params["Data"]["noisy_pos_variance"] rot_noise = params["Data"]["noisy_rot_variance"] noisy_poses = get_pose_noise_np(setup_params["max_envs"], setup_params["trajectory_length"], pos_noise, rot_noise) save_noisy_poses(noisy_poses) print("saved noisy poses for " + str(setup_params["max_envs"]) + " envs")
def count_avg_num_steps(): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] P.get_current_parameters()["Data"]["locking"] = False train_envs, dev_envs, test_envs = get_restricted_env_id_lists() train_envs = [e for e in train_envs if e >= 6000] count_avg_num_steps_on_data(train_envs)
def train_supervised(): initialize_experiment() setup = get_current_parameters()["Setup"] supervised_params = get_current_parameters()["Supervised"] num_epochs = supervised_params["num_epochs"] model, model_loaded = load_model() print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() filename = "supervised_" + setup["model"] + "_" + setup["run_name"] start_filename = "tmp/" + filename + "_epoch_" + str( supervised_params["start_epoch"]) if supervised_params["start_epoch"] > 0: if file_exists(start_filename, ""): load_pytorch_model(model, start_filename) else: print("Couldn't continue training. Model file doesn't exist at:") print(start_filename) exit(-1) trainer = Trainer(model, epoch=supervised_params["start_epoch"], name=setup["model"], run_name=setup["run_name"]) print("Beginning training...") best_test_loss = 1000 for epoch in range(num_epochs): train_loss = trainer.train_epoch(train_data=None, train_envs=train_envs, eval=False) trainer.model.correct_goals = 0 trainer.model.total_goals = 0 test_loss = trainer.train_epoch(train_data=None, train_envs=dev_envs, eval=True) print("GOALS: ", trainer.model.correct_goals, trainer.model.total_goals) if test_loss < best_test_loss: best_test_loss = test_loss save_pytorch_model(trainer.model, filename) print("Saved model in:", filename) print("Epoch", epoch, "train_loss:", train_loss, "test_loss:", test_loss) save_pytorch_model(trainer.model, "tmp/" + filename + "_epoch_" + str(epoch)) if hasattr(trainer.model, "save"): trainer.model.save(epoch) save_pretrained_weights(trainer.model, setup["run_name"])
def train_supervised_bidomain(): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] model_sim, _ = load_model(setup["model"], setup["sim_model_file"], domain="sim") model_real, _ = load_model(setup["model"], setup["real_model_file"], domain="real") model_critic, _ = load_model(setup["critic_model"], setup["critic_model_file"]) print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() env_list_name = setup.get("eval_env_set", "dev") if env_list_name == "dev": print("Using DEV envs") use_envs = dev_envs elif env_list_name == "train": print("Using TRAIN envs") use_envs = train_envs elif env_list_name == "test": print("Using TEST envs") use_envs = test_envs else: raise ValueError(f"Unknown env set {env_list_name}") env_range_start = setup.get("env_range_start") if env_range_start > 0: use_envs = [e for e in use_envs if e >= env_range_start] env_range_end = setup.get("env_range_end") if env_range_end > 0: use_envs = [e for e in use_envs if e < env_range_end] restricted_domain = "simulator" if restricted_domain == "simulator": # Load dummy model for real domain model_real, _ = load_model(setup["model"], setup["sim_model_file"], domain="sim") model_sim.set_save_path_overlays(True) elif restricted_domain == "real": # Load dummy model for sim domain model_sim, _ = load_model(setup["model"], setup["real_model_file"], domain="real") model_real.set_save_path_overlays(True) else: model_real.set_save_path_overlays(True) model_sim.set_save_path_overlays(True) trainer = TrainerBidomain(model_real, model_sim, model_critic, epoch=0) trainer.train_epoch(env_list=use_envs, eval=True, restricted_domain=restricted_domain) if restricted_domain != "simulator": model_real.print_metrics() if restricted_domain != "real": model_sim.print_metrics()
def sample_real_data_subset(): global env_list, num_env_groups P.initialize_experiment() if env_list == "DEV": train_i, dev_i, test_i = get_restricted_env_id_lists() env_list = dev_i elif env_list == "TEST": train_i, dev_i, test_i = get_restricted_env_id_lists() env_list = test_i # Each 5 subsequent environments are the same. First sample groups, then sample environments groups = set() for env in env_list: groups.add(int(env/5)) groups = list(groups) group_envs_rel = {} pick_groups = random.sample(groups, num_env_groups) for group in pick_groups: group_envs_rel[group] = [] i = 0 while i < envs_each_group: rint = random.randint(0,4) if rint not in group_envs_rel[group]: group_envs_rel[group].append(rint) i += 1 else: # Retry this loop iteration continue env_ids_out = [] for group, env_rels in group_envs_rel.items(): for env_rel in env_rels: env_id = group * 5 + env_rel env_ids_out.append(env_id) print(f"Sampled {len(env_ids_out)} envs:") print(list(sorted(env_ids_out)))
def collect_supervised_data(): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] train_envs, dev_envs, test_envs = get_restricted_env_id_lists() # if P.get_current_parameters()["Setup"].get("env_set") == "train": print("Collecting data for training envs") collect_data_on_env_list(train_envs) elif P.get_current_parameters()["Setup"].get("env_set") == "dev": print("Collecting data for dev envs") collect_data_on_env_list(dev_envs) else: print("Collecting data for both training and dev envs") collect_data_on_env_list(train_envs) collect_data_on_env_list(dev_envs)
def train_rl(): initialize_experiment() setup = get_current_parameters()["Setup"] params = get_current_parameters()["RL"] print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() filename = "rl_" + setup["model"] + "_" + setup["run_name"] trainer = TrainerRL(params=dict_merge(setup, params)) for start_epoch in range(10000): epfname = epoch_filename(filename, start_epoch) path = os.path.join(get_model_dir(), str(epfname) + ".pytorch") if not os.path.exists(path): break if start_epoch > 0: print(f"CONTINUING RL TRAINING FROM EPOCH: {start_epoch}") load_pytorch_model(trainer.full_model, epoch_filename(filename, start_epoch - 1)) trainer.set_start_epoch(start_epoch) print("Beginning training...") best_dev_reward = -1e+10 for epoch in range(start_epoch, 10000): train_reward, metrics = trainer.train_epoch(eval=False, envs="train") # TODO: Test on just a few dev environments # TODO: Take most likely or mean action when testing dev_reward, metrics = trainer.train_epoch(eval=True, envs="dev") #dev_reward, metrics = trainer.train_epoch(eval=True, envs="dev") dev_reward = 0 #if dev_reward >= best_dev_reward: # best_dev_reward = dev_reward # save_pytorch_model(trainer.full_model, filename) # print("Saved model in:", filename) print("Epoch", epoch, "train reward:", train_reward, "dev reward:", dev_reward) save_pytorch_model(trainer.full_model, epoch_filename(filename, epoch)) if hasattr(trainer.full_model, "save"): trainer.full_model.save(epoch)
def view_collected_data(args): train_envs, dev_envs, test_envs = get_restricted_env_id_lists( args.max_envs) total_len = 0 count = 0 maxlen = 0 for env_id in train_envs: print("Showing env id: ", env_id) data = load_single_env_supervised_data(env_id) for sample in data: Presenter().show_sample(sample["state"], sample["action"], 0, sample["instruction"]) print("Image size: ", sample["state"].image.shape) print("Pose: ", sample["state"].get_pos_3d(), sample["state"].get_rot_euler()) cv2.waitKey() total_len += len(data)
def copy_seed_dataset(from_dataset, to_dataset, seed_count): from_dir = get_dataset_dir(from_dataset) to_dir = get_dataset_dir(to_dataset) if os.path.exists(to_dir): print("DATASET EXISTS! Continue where left off?") c = input(" (y/n) >>> ") if c == "y": return else: raise ValueError( f"Not continuing: Dataset {to_dataset} exists. Delete it if you like and try again" ) os.makedirs(to_dir) from_files = os.listdir(from_dir) train_ids, dev_ids, test_ids = get_restricted_env_id_lists() train_ids = set(train_ids) dev_ids = set(dev_ids) test_ids = set(test_ids) file_envs = [ int(f.split("supervised_train_data_env_")[1]) for f in from_files ] files_and_envs = list(zip(from_files, file_envs)) random.shuffle(files_and_envs) files_to_copy = [] train_envs_copied = 0 for file, env in files_and_envs: if env in train_ids and train_envs_copied < seed_count: files_to_copy.append(file) if env in train_ids: train_envs_copied += 1 print( f"Copying {train_envs_copied} train envs, and all dev/test envs from {from_dataset} to {to_dataset}" ) for file in files_to_copy: from_path = os.path.join(from_dir, file) to_path = os.path.join(to_dir, file) shutil.copy(from_path, to_path)
def browse_pvn_dataset(): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] model_sim, _ = load_model(setup["model"], setup["sim_model_file"], domain="sim") data_params = P.get_current_parameters()["Training"] print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() #dom="real" dom = "sim" dataset = model_sim.get_dataset( data=None, envs=train_envs, domain=dom, dataset_names=data_params[f"{dom}_dataset_names"], dataset_prefix="supervised", eval=False, halfway_only=False) p = Presenter() for example in dataset: if example is None: continue md = example["md"][0] print( f"Showing example: {md['env_id']}:{md['set_idx']}:{md['seg_idx']}") print(f" instruction: {md['instruction']}") exec_len = len(example["images"]) for i in range(exec_len): print(f" timestep: {i}") img_i = example["images"][i] lm_fpv_i = example["lm_pos_fpv"][i] if lm_fpv_i is not None: img_i = p.plot_pts_on_torch_image(img_i, lm_fpv_i.long()) p.show_image(img_i, "fpv_img_i", scale=4, waitkey=True)
def sample_human_envs(): train_envs, dev_envs, test_envs = get_restricted_env_id_lists() random.shuffle(test_envs) human_envs = test_envs[:NUM_ENVS] human_envs = sorted(human_envs) save_json(human_envs, get_human_eval_envs_path())
c5 = (charray / 5).astype(np.int64) copies = [c5[c] in c5[:c] or c in c5[c+1:] for c in range(len(c5))] num_copies = np.asarray(copies).sum() print("num_envs_same: ", num_copies) if num_copies > 0: return False return True if __name__ == "__main__": P.initialize_experiment() real_data_dir = get_dataset_dir("real") files = os.listdir(real_data_dir) available_env_ids = set([int(f.split("supervised_train_data_env_")[1]) for f in files]) train_ids, dev_ids, test_ids = get_restricted_env_id_lists() train_ids = set(train_ids) dev_ids = set(dev_ids) test_ids = set(test_ids) avail_train_ids = list(train_ids.intersection(available_env_ids)) avail_dev_ids = list(dev_ids.intersection(available_env_ids)) avail_test_ids = list(test_ids.intersection(available_env_ids)) print(f"Making subsets from total envs: {len(avail_train_ids)}") splits_out = {} choice = [] prev_split = splits[0] while not first_choice_ok(choice):
def train_top_down_pred(args, max_epoch=SUPERVISED_EPOCHS): initialize_experiment(args.run_name, args.setup_name) model, model_loaded = load_model() # TODO: Get batch size from global parameter server when it exists batch_size = 1 if \ args.model == "top_down" or \ args.model == "top_down_prior" or \ args.model == "top_down_sm" or \ args.model == "top_down_pretrain" or \ args.model == "top_down_goal_pretrain" or \ args.model == "top_down_nav" or \ args.model == "top_down_cond" \ else BATCH_SIZE lr = 0.001 # * batch_size trainer = Trainer(model, epoch=args.start_epoch, name=args.model, run_name=args.run_name) train_envs, dev_envs, test_envs = get_restricted_env_id_lists( max_envs=args.max_envs) filename = "top_down_" + args.model + "_" + args.run_name if args.restore_weights_name is not None: restore_pretrained_weights(model, args.restore_weights_name, args.fix_restored_weights) print("Beginning training...") best_test_loss = 1000 validation_loss = [] for epoch in range(SUPERVISED_EPOCHS): train_loss = -1 if not args.eval_pretrain: train_loss = trainer.train_epoch(train_envs=train_envs, eval=False) test_loss = trainer.train_epoch(train_envs=dev_envs, eval=True) validation_loss.append([epoch, test_loss]) if not args.eval_pretrain: if test_loss < best_test_loss: best_test_loss = test_loss save_pytorch_model(trainer.model, filename) print("Saved model in:", filename) print("Epoch", epoch, "train_loss:", train_loss, "test_loss:", test_loss) save_pytorch_model(trainer.model, "tmp/" + filename + "_epoch_" + str(epoch)) save_pretrained_weights(trainer.model, args.run_name) else: break if max_epoch is not None and epoch > max_epoch: print("Reached epoch limit!") break test_loss_dir = get_model_dir( ) + "/test_loss/" + filename + "_test_loss.csv" validation_loss = pd.DataFrame(validation_loss, columns=['epoch', "test_loss"]) validation_loss.to_csv(test_loss_dir, index=False)
from data_io.instructions import get_restricted_env_id_lists, get_all_instructions import parameters.parameter_server as P if __name__ == "__main__": P.initialize_experiment() train_envs, dev_envs, test_envs = get_restricted_env_id_lists() train_i, dev_i, test_i, corpus = get_all_instructions() train_i_envs = set([int(i) for i in train_i.keys()]) for test_env in test_envs: assert test_env not in train_i_envs, "FAIL" for dev_env in test_envs: assert dev_env not in train_i_envs, "FAIL" print("OK")
def train_supervised_worker(rl_process_conn): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] rlsup = P.get_current_parameters()["RLSUP"] setup["trajectory_length"] = setup["sup_trajectory_length"] run_name = setup["run_name"] supervised_params = P.get_current_parameters()["Supervised"] num_epochs = supervised_params["num_epochs"] sup_device = rlsup.get("sup_device", "cuda:1") model_oracle_critic = None print("SUPP: Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() # Load the starter model and save it at epoch 0 # Supervised worker to use GPU 1, RL will use GPU 0. Simulators run on GPU 2 model_sim = load_model(setup["sup_model"], setup["sim_model_file"], domain="sim")[0].to(sup_device) model_real = load_model(setup["sup_model"], setup["real_model_file"], domain="real")[0].to(sup_device) model_critic = load_model(setup["sup_critic_model"], setup["critic_model_file"])[0].to(sup_device) # ---------------------------------------------------------------------------------------------------------------- print("SUPP: Initializing trainer") rlsup_params = P.get_current_parameters()["RLSUP"] sim_seed_dataset = rlsup_params.get("sim_seed_dataset") # TODO: Figure if 6000 or 7000 here trainer = TrainerBidomainBidata(model_real, model_sim, model_critic, model_oracle_critic, epoch=0) train_envs_common = [e for e in train_envs if 6000 <= e < 7000] train_envs_sim = [e for e in train_envs if e < 7000] dev_envs_common = [e for e in dev_envs if 6000 <= e < 7000] dev_envs_sim = [e for e in dev_envs if e < 7000] sim_datasets = [rl_dataset_name(run_name)] real_datasets = ["real"] trainer.set_dataset_names(sim_datasets=sim_datasets, real_datasets=real_datasets) # ---------------------------------------------------------------------------------------------------------------- for start_sup_epoch in range(10000): epfname = epoch_sup_filename(run_name, start_sup_epoch, model="stage1", domain="sim") path = os.path.join(get_model_dir(), str(epfname) + ".pytorch") if not os.path.exists(path): break if start_sup_epoch > 0: print(f"SUPP: CONTINUING SUP TRAINING FROM EPOCH: {start_sup_epoch}") load_pytorch_model( model_real, epoch_sup_filename(run_name, start_sup_epoch - 1, model="stage1", domain="real")) load_pytorch_model( model_sim, epoch_sup_filename(run_name, start_sup_epoch - 1, model="stage1", domain="sim")) load_pytorch_model( model_critic, epoch_sup_filename(run_name, start_sup_epoch - 1, model="critic", domain="critic")) trainer.set_start_epoch(start_sup_epoch) # ---------------------------------------------------------------------------------------------------------------- print("SUPP: Beginning training...") for epoch in range(start_sup_epoch, num_epochs): # Tell the RL process that a new Stage 1 model is ready for loading print("SUPP: Sending model to RL") model_sim.reset() rl_process_conn.send( ["stage1_model_state_dict", model_sim.state_dict()]) if DEBUG_RL: while True: sleep(1) if not sim_seed_dataset: ddir = get_dataset_dir(rl_dataset_name(run_name)) os.makedirs(ddir, exist_ok=True) while len(os.listdir(ddir)) < 20: print("SUPP: Waiting for rollouts to appear") sleep(3) print("SUPP: Beginning Epoch") train_loss = trainer.train_epoch(env_list_common=train_envs_common, env_list_sim=train_envs_sim, eval=False) test_loss = trainer.train_epoch(env_list_common=dev_envs_common, env_list_sim=dev_envs_sim, eval=True) print("SUPP: Epoch", epoch, "train_loss:", train_loss, "test_loss:", test_loss) save_pytorch_model( model_real, epoch_sup_filename(run_name, epoch, model="stage1", domain="real")) save_pytorch_model( model_sim, epoch_sup_filename(run_name, epoch, model="stage1", domain="sim")) save_pytorch_model( model_critic, epoch_sup_filename(run_name, epoch, model="critic", domain="critic"))
def generate_rollout_debug_visualizations(): setup = P.get_current_parameters()["Setup"] dataset_name = setup.get("viz_dataset_name") or get_eval_tmp_dataset_name(setup["model"], setup["run_name"]) domain = setup.get("viz_domain") or ("real" if setup.get("real_drone") else "sim") run_name = setup.get("original_run_name") or setup.get("run_name") specific_envs = setup.get("only_specific_envs") # For collecting information for visualization examples specific_segments = [ # running example (6827, 0, 4), # successful examples (6169, 0, 9), (6825, 0, 8), (6857, 0, 9), # failure examples (6169, 0, 2), (6299, 0, 9), (6634, 0, 8), (6856, 0, 9), (6857, 0, 8), ] specific_segments += [ # good sim, lousy real (6419, 0, 5), (6569, 0, 6), (6634, 0, 6), (6917, 0, 7), ] specific_envs = [s[0] for s in specific_segments] # Generate all #specific_envs = list(range(6000, 7000, 1)) #specific_segments = None # Some quick params. TODO: Bring this into json viz_params = { "ego_vdist": False, "draw_landmarks": False, "draw_topdown": True, "draw_drone": True, "draw_trajectory": True, "draw_fov": False, "include_vdist": False, "include_layer": None, "include_instr": False } print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() # TODO: Grab the correct env list env_list = dev_envs viz = RolloutVisualizer(resolution=576) base_dir = os.path.join(get_rollout_debug_viz_dir(), f"{dataset_name}-{domain}") os.makedirs(base_dir, exist_ok=True) for env_id in env_list: if specific_envs and env_id not in specific_envs: print("Skipping", env_id) continue try: env_data = load_single_env_from_dataset(dataset_name, env_id, "supervised") except FileNotFoundError as e: print(f"Skipping env: {env_id}") continue if len(env_data) == 0: print(f"Skipping env: {env_id}. Rollout exists but is EMPTY!") continue segs = split_into_segs(env_data) for seg in segs: lag_start = 1.5 end_lag = 1.5 seg_idx = seg[0]["seg_idx"] if specific_segments and (env_id, 0, seg_idx) not in specific_segments: continue seg_name = f"{env_id}:0:{seg_idx}-{domain}" gif_filename = f"{seg_name}-roll" instr_filename = f"{seg_name}-instr.txt" this_dir = os.path.join(base_dir, gif_filename) os.makedirs(this_dir, exist_ok=True) base_path = os.path.join(this_dir, gif_filename) if os.path.exists(os.path.join(this_dir, instr_filename)): continue # Animation with just the drone frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, viz_params) save_frames(viz, frames, f"{base_path}-exec", fps=5.0, start_lag=lag_start, end_lag=end_lag, formats=Y_FMT) # Save instructionto with open(os.path.join(this_dir, instr_filename), "w") as fp: fp.write(seg[0]["instruction"]) # Animation of action frames = viz.action_visualization(env_id, seg_idx, seg, domain, "action") save_frames(viz, frames, f"{base_path}-action", fps=5.0, start_lag=lag_start, end_lag=end_lag) # Animation of actions #action_frames = viz.grab_frames(env_id, seg_idx, seg, domain, "action", scale=4) #save_frames(viz, action_frames, f"{base_path}-action", fps=5.0, start_lag=lag_start, end_lag=end_lag) # Generate and save gif # Bare top-down view mod_params = deepcopy(viz_params) mod_params["draw_drone"] = False mod_params["draw_trajectory"] = False frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, mod_params) save_frames(viz, frames, f"{base_path}-top-down", fps=5.0, start_lag=lag_start, end_lag=end_lag) mod_params["draw_drone"] = True mod_params["draw_trajectory"] = False frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, mod_params) save_frames(viz, frames, f"{base_path}-top-down-drn", fps=5.0, start_lag=lag_start, end_lag=end_lag) # Egocentric visitation distributions vdist_r_frames = viz.grab_frames(env_id, seg_idx, seg, domain, "v_dist_r_inner") save_frames(viz, vdist_r_frames, f"{base_path}-ego-vdist", fps=5.0, start_lag=lag_start, end_lag=end_lag) # Map struct map_struct_frames = viz.grab_frames(env_id, seg_idx, seg, domain, "map_struct") save_frames(viz, map_struct_frames, f"{base_path}-ego-map-struct", fps=5.0, start_lag=lag_start, end_lag=end_lag) # Egocentric observation mask ego_obs_mask_frames = viz.grab_frames(env_id, seg_idx, seg, domain, "ego_obs_mask") save_frames(viz, ego_obs_mask_frames, f"{base_path}-ego-obs-mask", fps=5.0, start_lag=lag_start, end_lag=end_lag) def save_map_permutations(file_prefix, incl_layer): mod_params = deepcopy(viz_params) if incl_layer == "vdist": mod_params["include_vdist"] = True else: mod_params["include_layer"] = incl_layer print(f"GENERATING: {file_prefix}") # Non-overlaid, without trajectory mod_params["draw_drone"] = False mod_params["draw_topdown"] = False mod_params["draw_trajectory"] = False frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, mod_params) save_frames(viz, frames, f"{file_prefix}", fps=5.0, start_lag=lag_start, end_lag=end_lag, formats=Y_FMT) print(f"GENERATING: {file_prefix}-ov") # Overlaid, without trajectory mod_params["draw_topdown"] = True frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, mod_params) save_frames(viz, frames, f"{file_prefix}-ov", fps=5.0, start_lag=lag_start, end_lag=end_lag, formats=D_FMT) print(f"GENERATING: {file_prefix}-ov-path") # Overlaid, with trajectory mod_params["draw_drone"] = True mod_params["draw_trajectory"] = True frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, mod_params) save_frames(viz, frames, f"{file_prefix}-ov-path", fps=5.0, start_lag=lag_start, end_lag=end_lag, formats=Y_FMT) print(f"GENERATING: {file_prefix}-path") # Non-overlaid, with trajectory mod_params["draw_topdown"] = False mod_params["draw_drone"] = True mod_params["draw_trajectory"] = True frames = viz.top_down_visualization(env_id, seg_idx, seg, domain, mod_params) save_frames(viz, frames, f"{file_prefix}-path", fps=5.0, start_lag=lag_start, end_lag=end_lag, formats=D_FMT) save_map_permutations(f"{base_path}-vdist", "vdist") save_map_permutations(f"{base_path}-semantic-map", "S_W") save_map_permutations(f"{base_path}-semantic-map-gray", "S_W_Gray") save_map_permutations(f"{base_path}-proj-features", "F_W") save_map_permutations(f"{base_path}-grounding-map", "R_W") save_map_permutations(f"{base_path}-grounding-map-gray", "R_W_Gray") save_map_permutations(f"{base_path}-mask", "M_W") save_map_permutations(f"{base_path}-accum-mask", "M_W_accum") save_map_permutations(f"{base_path}-accum-mask-inv", "M_W_accum_inv") # Animation of FPV features fpv_feature_frames = viz.grab_frames(env_id, seg_idx, seg, domain, "F_C") save_frames(viz, fpv_feature_frames, f"{base_path}-features-fpv", fps=5.0, start_lag=lag_start, end_lag=end_lag) # Animation of FPV images fpv_image_frames = viz.grab_frames(env_id, seg_idx, seg, domain, "image", scale=4) save_frames(viz, fpv_image_frames, f"{base_path}-image", fps=5.0, start_lag=lag_start, end_lag=end_lag) frames = viz.overlay_frames(fpv_image_frames, fpv_feature_frames) save_frames(viz, frames, f"{base_path}-features-fpv-ov", fps=5.0, start_lag=lag_start, end_lag=end_lag) num_frames = len(frames) # Clip rollout videos to correct rollout duration and re-save rollout_dir = get_rollout_video_dir(run_name=run_name) if os.path.isdir(rollout_dir): print("Processing rollout videos") actual_rollout_duration = num_frames / 5.0 ceiling_clip = viz.load_video_clip(env_id, seg_idx, seg, domain, "ceiling", rollout_dir) duration_with_lag = lag_start + actual_rollout_duration + end_lag try: if ceiling_clip is not None: if ceiling_clip.duration > duration_with_lag: start = ceiling_clip.duration - end_lag - duration_with_lag ceiling_clip = ceiling_clip.cutout(0, start) #ceiling_clip = ceiling_clip.cutout(duration_with_lag, ceiling_clip.duration) save_frames(viz, ceiling_clip, f"{base_path}-ceiing_cam-clipped", fps=ceiling_clip.fps) corner_clip = viz.load_video_clip(env_id, seg_idx, seg, domain, "corner", rollout_dir) if corner_clip is not None: if corner_clip.duration > actual_rollout_duration + end_lag: start = corner_clip.duration - end_lag - duration_with_lag corner_clip = corner_clip.cutout(0, start) #corner_clip = corner_clip.cutout(duration_with_lag, corner_clip.duration) save_frames(viz, corner_clip, f"{base_path}-corner_cam-clipped", fps=corner_clip.fps) except Exception as e: print("Video encoding error! Copying manually") print(e) try: in_ceil_file = os.path.join(rollout_dir, f"rollout_ceiling_{env_id}-0-{seg_idx}.mkv") in_corn_file = os.path.join(rollout_dir, f"rollout_corner_{env_id}-0-{seg_idx}.mkv") out_ceil_file = f"{base_path}-ceiling_cam-full.mkv" out_corn_file = f"{base_path}-corner_cam-full.mkv" shutil.copy(in_ceil_file, out_ceil_file) shutil.copy(in_corn_file, out_corn_file) except Exception as e: print("Failed copying videos! SKipping") print("ding")
def train_dagger_simple(): # ---------------------------------------------------------------------------------------------------------------- # Load params and configure stuff P.initialize_experiment() params = P.get_current_parameters()["SimpleDagger"] setup = P.get_current_parameters()["Setup"] num_iterations = params["num_iterations"] sim_seed_dataset = params.get("sim_seed_dataset") run_name = setup["run_name"] device = params.get("device", "cuda:1") dataset_limit = params.get("dataset_size_limit_envs") seed_count = params.get("seed_count") # Trigger rebuild if necessary before going into all the threads and processes _ = get_restricted_env_id_lists(full=True) # Initialize the dataset if sim_seed_dataset: copy_seed_dataset(from_dataset=sim_seed_dataset, to_dataset=dagger_dataset_name(run_name), seed_count=seed_count or dataset_limit) gap = 0 else: # TODO: Refactor this into a prompt function data_path = get_dataset_dir(dagger_dataset_name(run_name)) if os.path.exists(data_path): print("DATASET EXISTS! Continue where left off?") c = input(" (y/n) >>> ") if c != "y": raise ValueError( f"Not continuing: Dataset {data_path} exists. Delete it if you like and try again" ) else: os.makedirs(data_path, exist_ok=True) gap = dataset_limit - len(os.listdir(data_path)) print("SUPP: Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() # ---------------------------------------------------------------------------------------------------------------- # Load / initialize model model = load_model(setup["model"], setup["model_file"], domain="sim")[0].to(device) oracle = load_model("oracle")[0] # ---------------------------------------------------------------------------------------------------------------- # Continue where we left off - load the model and set the iteration/epoch number for start_iteration in range(10000): epfname = epoch_dag_filename(run_name, start_iteration) path = os.path.join(get_model_dir(), str(epfname) + ".pytorch") if not os.path.exists(path): break if start_iteration > 0: print( f"DAG: CONTINUING DAGGER TRAINING FROM ITERATION: {start_iteration}" ) load_pytorch_model(model, epoch_dag_filename(run_name, start_iteration - 1)) # ---------------------------------------------------------------------------------------------------------------- # Intialize trainer trainer = Trainer(model, epoch=start_iteration, name=setup["model"], run_name=setup["run_name"]) trainer.set_dataset_names([dagger_dataset_name(run_name)]) # ---------------------------------------------------------------------------------------------------------------- # Initialize policy roller roller = SimpleParallelPolicyRoller( num_workers=params["num_workers"], device=params["device"], policy_name=setup["model"], policy_file=setup["model_file"], oracle=oracle, dataset_save_name=dagger_dataset_name(run_name), no_reward=True) rollout_sampler = RolloutSampler(roller) # ---------------------------------------------------------------------------------------------------------------- # Train DAgger - loop over iteartions, in each, prune, rollout and train an epoch print("SUPP: Beginning training...") for iteration in range(start_iteration, num_iterations): print(f"DAG: Starting iteration {iteration}") # Remove extra rollouts to keep within DAggerFM limit prune_dataset(run_name, dataset_limit) # Rollout and collect more data for training and evaluation policy_state = model.get_policy_state() rollout_sampler.sample_n_rollouts( n=gap if iteration == 0 else params["train_envs_per_iteration"], policy_state=policy_state, sample=False, envs="train", dagger_beta=dagger_beta(params, iteration)) eval_rollouts = rollout_sampler.sample_n_rollouts( n=params["eval_envs_per_iteration"], policy_state=policy_state, sample=False, envs="dev", dagger_beta=0) # Kill airsim instances so that they don't take up GPU memory and in general slow things down during training roller.kill_airsim() # Evaluate success / metrics and save to tensorboard if setup["eval_nl"]: evaler = DataEvalNL(run_name, entire_trajectory=False, save_images=False) evaler.evaluate_dataset(eval_rollouts) results = evaler.get_results() print("Results:", results) evaler.write_summaries(setup["run_name"], "dagger_eval", iteration) # Do one epoch of supervised training print("SUPP: Beginning Epoch") train_loss = trainer.train_epoch(train_envs=train_envs, eval=False) #test_loss = trainer.train_epoch(env_list_common=dev_envs_common, env_list_sim=dev_envs_sim, eval=True) # Save the model to file print("SUPP: Epoch", iteration, "train_loss:", train_loss) save_pytorch_model(model, epoch_dag_filename(run_name, iteration))
def train_supervised_bidomain(): P.initialize_experiment() setup = P.get_current_parameters()["Setup"] supervised_params = P.get_current_parameters()["Supervised"] num_epochs = supervised_params["num_epochs"] model_sim, _ = load_model(setup["model"], setup["sim_model_file"], domain="sim") model_real, _ = load_model(setup["model"], setup["real_model_file"], domain="real") model_critic, _ = load_model(setup["critic_model"], setup["critic_model_file"]) if P.get_current_parameters()["Training"].get("use_oracle_critic", False): model_oracle_critic, _ = load_model(setup["critic_model"], setup["critic_model_file"]) # This changes the name in the summary writer to get a different color plot oname = model_oracle_critic.model_name model_oracle_critic.set_model_name(oname + "_oracle") model_oracle_critic.model_name = oname else: model_oracle_critic = None print("Loading data") train_envs, dev_envs, test_envs = get_restricted_env_id_lists() real_filename = f"supervised_{setup['model']}_{setup['run_name']}_real" sim_filename = f"supervised_{setup['model']}_{setup['run_name']}_sim" critic_filename = f"supervised_{setup['critic_model']}_{setup['run_name']}_critic" # TODO: (Maybe) Implement continuing of training # Bidata means that we treat Lani++ and LaniOriginal examples differently, only computing domain-adversarial stuff on Lani++ bidata = P.get_current_parameters()["Training"].get("bidata", False) if bidata == "v2": trainer = TrainerBidomainBidata(model_real, model_sim, model_critic, model_oracle_critic, epoch=0) train_envs_common = [e for e in train_envs if 6000 <= e < 7000] train_envs_sim = train_envs dev_envs_common = [e for e in dev_envs if 6000 <= e < 7000] dev_envs_sim = dev_envs elif bidata: trainer = TrainerBidomainBidata(model_real, model_sim, model_critic, model_oracle_critic, epoch=0) train_envs_common = [e for e in train_envs if 6000 <= e < 7000] train_envs_sim = [e for e in train_envs if e < 6000] dev_envs_common = [e for e in dev_envs if 6000 <= e < 7000] dev_envs_sim = [e for e in dev_envs if e < 6000] else: trainer = TrainerBidomain(model_real, model_sim, model_critic, model_oracle_critic, epoch=0) print("Beginning training...") best_test_loss = 1000 for epoch in range(num_epochs): if bidata: train_loss = trainer.train_epoch(env_list_common=train_envs_common, env_list_sim=train_envs_sim, eval=False) test_loss = trainer.train_epoch(env_list_common=dev_envs_common, env_list_sim=dev_envs_sim, eval=True) else: train_loss = trainer.train_epoch(env_list=train_envs, eval=False) test_loss = trainer.train_epoch(env_list=dev_envs, eval=True) if test_loss < best_test_loss: best_test_loss = test_loss save_pytorch_model(model_real, real_filename) save_pytorch_model(model_sim, sim_filename) save_pytorch_model(model_critic, critic_filename) print(f"Saved models in: \n Real: {real_filename} \n Sim: {sim_filename} \n Critic: {critic_filename}") print ("Epoch", epoch, "train_loss:", train_loss, "test_loss:", test_loss) save_pytorch_model(model_real, f"tmp/{real_filename}_epoch_{epoch}") save_pytorch_model(model_sim, f"tmp/{sim_filename}_epoch_{epoch}") save_pytorch_model(model_critic, f"tmp/{critic_filename}_epoch_{epoch}")