def main(system_env, system, setup, ep, from_exported, network_type, outputfn):
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=state_size[system], encoder=mpnet.encoder.cuda()).cuda()
else:
if system == 'quadrotor_obs':
from networks.mpnet_quadrotor_obs import MPNet
elif system == 'cartpole_obs':
from networks.mpnet_cartpole_obs import MPNet
elif system == 'acrobot_obs':
from networks.mpnet_acrobot_obs import MPNet
elif system == 'car_obs':
from networks.mpnet_car_obs import MPNet
mpnet = MPNet(
ae_input_size=32,
ae_output_size=output_size[system],
in_channels=in_channel[system],
state_size=state_size[system]).cuda()
costnet = CostNet(ae_input_size=32, ae_output_size=output_size[system], in_channels=in_channel[system], state_size=state_size[system], encoder=mpnet.encoder).cuda()
costnet.load_state_dict(torch.load('output/{}/{}/{}/ep{}.pth'.format(system, setup, network_type, ep)))
costnet.eval()
# costnet.mlp.eval()
# costnet.mlp.dropout = False
Path("exported/output/{}".format(system)).mkdir(exist_ok=True)
export(costnet, setup=setup, system_env=system_env, system=system, exported_path="exported/output/{}/{}".format(system, outputfn))
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)