def main(args):
if args.dt is not None:
if args.abr == 'pensieve':
pensiedt.PensieveDT().main(args)
elif args.abr == 'robustmpc':
robustmdt.RobustMPCDT().main(args)
elif args.abr == 'hotdash':
hotdadt.HotdashDT().main(args)
else:
if args.abr == 'pensieve':
pensieve.Pensieve().main(args)
elif args.abr == 'robustmpc':
robustmpc.RobustMPC().main(args)
elif args.abr == 'hotdash':
hotdash.Hotdash().main(args)
all_cooked_bw=all_cooked_bw,
all_file_names=all_file_names)
time_calc = np.zeros((max_iters, 3))
if args.abr == 'pensieve':
teacher = pensieve.Pensieve()
student = pensiedt.PensieveDT()
predict = teacher.predict
elif args.abr == 'robustmpc':
teacher = robustmpc.RobustMPC()
student = robustmdt.RobustMPCDT()
predict = teacher.predict
elif args.abr == 'hotdash':
teacher = hotdash.Hotdash()
student = hotdadt.HotdashDT()
predict = teacher.predict
else:
raise NotImplementedError
t1 = time.time()
# Step 1: Initialization for the first iteration
trace = get_rollouts(env=net_env,
policy=teacher,
args=args,
n_batch_rollouts=n_batch_rollouts)
states.extend((state for state, _, _ in trace))
actions.extend((action for _, action, _ in trace))
serials.extend(serial for _, _, serial in trace)