def train(lr, epochs, batch, system, model, setup):
mpnet = MPNet(ae_input_size=32,
ae_output_size=128,
in_channels=1,
state_size=4,
control_size=2)
env_vox = torch.from_numpy(
np.load('{}/{}_env_vox.npy'.format(system, model))).float()
if torch.cuda.is_available():
mpnet = mpnet.cuda()
env_vox = env_vox.cuda()
#optimizer = torch.optim.Adagrad(mpnet.parameters(), lr=lr)
#optimizer = torch.optim.SGD(mpnet.parameters(), lr=lr, momentum=0.9)
optimizer = torch.optim.Adam(mpnet.parameters(), lr=lr)
#scheduler = StepLR(optimizer, step_size=10, gamma=0.9)
train_loader, test_loader = get_loader(system,
model,
batch_size=batch,
setup=setup)
for i in tqdm(range(epochs + 1)):
print("epoch {}".format(i))
train_loss = []
mpnet.train()
for data, label in tqdm(train_loader):
## prepare data
if torch.cuda.is_available():
data = data.cuda()
label = label.cuda()
inputs = data[:, 1:]
envs = env_vox[(data[:, 0]).long()]
## execute
optimizer.zero_grad()
output = mpnet(inputs, envs)
loss = F.mse_loss(output, label)
loss.backward()
optimizer.step()
train_loss.append(loss.item())
# break
#scheduler.step(i)
print("train loss:{}".format(np.mean(train_loss)))
for data, label in test_loader:
mpnet.eval()
eval_loss = []
with torch.no_grad():
if torch.cuda.is_available():
data = data.cuda()
label = label.cuda()
inputs = data[:, 1:]
envs = env_vox[(data[:, 0]).long()]
output = mpnet(inputs, envs)
loss = F.mse_loss(output, label)
eval_loss.append(loss.item())
print("eval loss:{}".format(np.mean(eval_loss)))
if i % 5 == 0:
Path("output/{}/{}".format(model, setup)).mkdir(parents=True,
exist_ok=True)
torch.save(mpnet.state_dict(),
"output/{}/{}/ep{}".format(model, setup, i))
def main(system_env, system, setup, ep):
mpnet = MPNet(
ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=state_size[system]).cuda()
costnet = CostNet(ae_input_size=32, ae_output_size=1024, in_channels=1, state_size=4, encoder=mpnet.encoder).cuda()
costnet.load_state_dict(torch.load('output/{}/{}/cost_so_far/ep{}.pth'.format(system, setup, ep)))
costnet.eval()
export(costnet, setup=setup, system_env=system_env, system=system, exported_path="exported/output/{}/cost_so_far_10k.pt".format(system))
def main(system_env, system, setup, ep, from_exported):
if from_exported:
from exported.export_mpnet_external_small_model import KMPNet, load_func, Encoder, MLP
mpnet = KMPNet(total_input_size=8,
AE_input_size=32,
mlp_input_size=40,
output_size=4,
CAE=Encoder,
MLP=MLP,
loss_f=None).cuda()
load_func(
mpnet,
'/media/arclabdl1/HD1/YLmiao/results/KMPnet_res/cartpole_obs_4_lr0.010000_Adagrad_step_200/kmpnet_epoch_3150_direction_0_step_200.pkl'
)
costnet = CostNet(ae_input_size=32,
ae_output_size=32,
in_channels=1,
state_size=4,
encoder=mpnet.encoder.cuda()).cuda()
else:
mpnet = MPNet(ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=state_size[system]).cuda()
costnet = CostNet(ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=4,
encoder=mpnet.encoder).cuda()
costnet.load_state_dict(
torch.load('output/{}/{}/cost_transit/ep{}.pth'.format(
system, setup, ep)))
costnet.eval()
export(costnet,
setup=setup,
system_env=system_env,
system=system,
exported_path="exported/output/{}/cost_10k.pt".format(system))
def main(ae_output_size, state_size, lr, epochs, batch, system_env, system,
setup, loss_type, load_from, network_type, data_type, label_type,
from_exported):
if from_exported:
import sys
sys.path.append(
"/media/arclabdl1/HD1/Linjun/mpc-mpnet-py/mpnet/exported")
from exported.export_mpnet_external_small_model import KMPNet, load_func, Encoder, MLP
mpnet = KMPNet(total_input_size=8,
AE_input_size=32,
mlp_input_size=40,
output_size=4,
CAE=Encoder,
MLP=MLP,
loss_f=None)
load_func(
mpnet,
'/media/arclabdl1/HD1/YLmiao/results/KMPnet_res/cartpole_obs_4_lr0.010000_Adagrad_step_200/kmpnet_epoch_3150_direction_0_step_200.pkl'
)
costnet = CostNet(ae_input_size=32,
ae_output_size=32,
in_channels=1,
state_size=4,
encoder=mpnet.encoder)
else:
from networks.mpnet import MPNet
mpnet = MPNet(ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=4)
mpnet.load_state_dict(
torch.load('output/{}/{}/{}/ep10000.pth'.format(
system, setup, load_from)))
costnet = CostNet(ae_input_size=32,
ae_output_size=1024,
in_channels=1,
state_size=4,
encoder=mpnet.encoder)
# for param in costnet.encoder.parameters():
# param.requires_grad = False
data_loaders = get_loader_cost(system_env,
system,
batch_size=batch,
setup=setup,
label_type=label_type,
data_type=data_type)
train_network(network=costnet,
data_loaders=data_loaders,
network_name=network_type,
lr=lr,
epochs=epochs,
batch=batch,
system_env=system_env,
system=system,
setup=setup,
using_step_lr=True,
step_size=50,
gamma=0.9,
loss_type=loss_type,
weight_save_epochs=50)