def main():
# Parse command line arguments.
parser = argparse.ArgumentParser()
parser.add_argument('--disp', action='store_true')
parser.add_argument('--task', default='insertion')
parser.add_argument('--agent', default='transporter')
parser.add_argument('--n_demos', default=100, type=int)
args = parser.parse_args()
name = f'{args.task}-{args.agent}-{args.n_demos}-'
print(name)
# Load and print results to console.
path = os.path.join('.')
curve = []
for fname in sorted(os.listdir(path)):
if name in fname and '.pkl' in fname:
n_steps = int(fname[(fname.rfind('-') + 1):-4])
data = pickle.load(open(fname, 'rb'))
rewards = []
for reward, _ in data:
rewards.append(reward)
rewards = np.array(rewards) * 100
score = np.mean(rewards)
std = np.std(rewards)
print(f' {n_steps} steps:\t{score:.1f}%\t± {std:.1f}%')
curve.append((n_steps, score, std))
# Plot results over training steps.
title = f'{args.agent} on {args.task} w/ {args.n_demos} demos'
ylabel = 'Testing Task Success (%)'
xlabel = '# of Training Steps'
if args.disp:
logs = {}
curve = np.array(curve)
logs[name] = (curve[:, 0], curve[:, 1], curve[:, 2])
fname = f'{name}-plot.png'
utils.plot(fname, title, ylabel, xlabel, data=logs, ylim=[0, 1])
print(f'Done. Plot image saved to: {fname}')
FILE_HEAD = 'test_results_insertion' # Daniel: ADJUST IF NECESSARY.
LAST_FEW_RUNS = 5 # Transporters paper did 5
for name, file in files.items():
print(f'\nPlotting: {name}, {file}')
x, y, std = [0], [0], [0]
for order in range(4):
if 'crop_befe_q' in file:
file1, file2 = file.split()
fname = file1 + f'{10**order}-{train_seed}' + file2 + '.pkl'
else:
fname = file + f'{10**order}-{train_seed}.pkl'
fname = os.path.join(FILE_HEAD, fname)
print(f'fname: {fname}')
if os.path.isfile(fname):
x.append(len(x))
data = pickle.load(open(fname, 'rb'))
data = np.float32(data)
ix = np.sort(np.unique(data[:, 0]))
iy = np.float32([data[data[:, 0] == i, 1] for i in ix])
iy = iy[-LAST_FEW_RUNS:].reshape(-1) # see comments above
std.append(np.std(iy))
y.append(np.mean(iy))
print(f' {np.mean(iy):0.3f} +/- {np.std(iy):0.3f}, for {fname}')
logs[name] = (x, y, std)
xticks = ['0', '1', '10', '100', '1000']
fname = os.path.join(os.getcwd(), 'figures',
f'plot_train_seed_{train_seed}.png')
utils.plot(fname, title, ylabel, xlabel, data=logs, xticks=xticks, ylim=[0, 1])
print(f'Done. Plot image saved to: {fname}')