def save_pytorch_model(model, name):
path = os.path.join(get_model_dir(), str(name) + ".pytorch")
dirname = os.path.dirname(path)
if not os.path.isdir(dirname):
try:
os.makedirs(dirname)
except Exception:
pass
torch.save(model.state_dict(), path)
def load_pytorch_model(model, name):
path = os.path.join(get_model_dir(), str(name) + ".pytorch")
print("path:", path)
model_dict = torch.load(path, map_location=lambda storage, loc: storage)
# Fix for loading PyTorch 0.3 models in PyTorch 0.4
model_dict_clone = model_dict.copy() # We can't mutate while iterating
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
model.load_state_dict(model_dict)
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 load_pytorch_model(model, name, pytorch3to4=False, namespace=None):
path = os.path.join(get_model_dir(), str(name) + ".pytorch")
model_dict = torch.load(path, map_location=lambda storage, loc: storage)
# Fix for loading PyTorch 0.3 models in PyTorch 0.4
if pytorch3to4:
model_dict_clone = model_dict.copy() # We can't mutate while iterating
for key, value in model_dict_clone.items():
if key.endswith(('running_mean', 'running_var')):
del model_dict[key]
if namespace is not None:
new_dict = {}
for k, v in model_dict.items():
if k.startswith(namespace):
k = k[len(namespace) + 1:]
new_dict[k] = v
model_dict = new_dict
if model is not None:
model.load_state_dict(model_dict)
else:
return model_dict
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_all_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)
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 train_rl_worker(sup_process_conn):
P.initialize_experiment()
setup = P.get_current_parameters()["Setup"]
setup["trajectory_length"] = setup["rl_trajectory_length"]
run_name = setup["run_name"]
rlsup = P.get_current_parameters()["RLSUP"]
params = P.get_current_parameters()["RL"]
num_rl_epochs = params["num_epochs"]
# These need to be distinguished between supervised and RL because supervised trains on ALL envs, RL only on 6000-7000
setup["env_range_start"] = setup["rl_env_range_start"]
setup["env_range_end"] = setup["rl_env_range_end"]
rl_device = rlsup.get("rl_device", "cuda:0")
trainer = TrainerRL(params=dict_merge(setup, params),
save_rollouts_to_dataset=rl_dataset_name(run_name),
device=rl_device)
# -------------------------------------------------------------------------------------
# TODO: Continue (including figure out how to initialize Supervised Stage 1 real/sim/critic and RL Stage 2 policy
start_rl_epoch = 0
for start_rl_epoch in range(num_rl_epochs):
epfname = epoch_rl_filename(run_name, start_rl_epoch, model="full")
path = os.path.join(get_model_dir(), str(epfname) + ".pytorch")
if not os.path.exists(path):
break
if start_rl_epoch > 0:
print(f"RLP: CONTINUING RL TRAINING FROM EPOCH: {start_rl_epoch}")
load_pytorch_model(
trainer.full_model,
epoch_rl_filename(run_name, start_rl_epoch - 1, model="full"))
trainer.set_start_epoch(start_rl_epoch)
# Wait for supervised process to send it's model
sleep(2)
# -------------------------------------------------------------------------------------
print("RLP: Beginning training...")
for rl_epoch in range(start_rl_epoch, num_rl_epochs):
# Get the latest Stage 1 model. Halt on the first epoch so that we can actually initialize the Stage 1
new_stage1_model_state_dict = receive_stage1_state(
sup_process_conn, halt=(rl_epoch == start_rl_epoch))
if new_stage1_model_state_dict:
print(f"RLP: Re-loading latest Stage 1 model")
trainer.reload_stage1(new_stage1_model_state_dict)
train_reward, metrics = trainer.train_epoch(epoch_num=rl_epoch,
eval=False,
envs="train")
dev_reward, metrics = trainer.train_epoch(epoch_num=rl_epoch,
eval=True,
envs="dev")
print("RLP: RL Epoch", rl_epoch, "train reward:", train_reward,
"dev reward:", dev_reward)
save_pytorch_model(trainer.full_model,
epoch_rl_filename(run_name, rl_epoch, model="full"))
save_pytorch_model(
trainer.full_model.stage1_visitation_prediction,
epoch_rl_filename(run_name, rl_epoch, model="stage1"))
save_pytorch_model(
trainer.full_model.stage2_action_generation,
epoch_rl_filename(run_name, rl_epoch, model="stage2"))
def file_exists(name):
full_path = os.path.join(get_dataset_dir(), str(name))
data_exists = os.path.exists(full_path)
full_path = os.path.join(get_model_dir(), str(name))
model_exists = os.path.exists(full_path)
return data_exists or model_exists
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))
