return mode
def get_n_actions(folder):
if 'rope' in folder:
n_actions = 20
elif 'cloth' in folder:
n_actions = 40
else:
raise Exception(folder)
return n_actions
if __name__ == '__main__':
n_gpus = torch.cuda.device_count()
folders = sys.argv[1:]
exps = [dict(folder=folders, mode=[get_mode(f) for f in folders],
n_actions=[get_n_actions(f) for f in folders]),
dict(goal_type=['random', 'flat']), dict(n_trials=[1000])]
exps = construct_variants(exps, name_key=None)
chunk_size = math.ceil(len(exps) / n_gpus)
worker_args = []
for i in range(n_gpus):
start, end = chunk_size * i, min(chunk_size * (i + 1), len(exps))
worker_args.append((i, exps[start:end]))
workers = [mp.Process(target=worker, args=arg) for arg in worker_args]
[w.start() for w in workers]
[w.join() for w in workers]
exps = [{
'root': ['data/rope'],
'state-size': [8],
'chunk-size': [10],
'overshooting-distance': [10],
'overshooting-kl-beta': [0],
'global-kl-beta': [0],
'trans_type': ['mlp'],
'embedding-size': [512],
'hidden-size': [128]
}, {
'seed': [0, 1, 2, 3]
}]
exps = construct_variants(exps, name_key='id')
if args.end == -1:
args.end = len(exps)
print(f'Running {args.end - args.start} / {len(exps)} experiments')
exps = exps[args.start:args.end]
n_exps = len(exps)
chunk_size = math.ceil(n_exps / args.n_gpus)
worker_args = []
for i in range(args.n_gpus):
start, end = chunk_size * i, min(chunk_size * (i + 1), n_exps)
worker_args.append((i, args.max_per_gpu, exps[start:end]))
workers = [mp.Process(target=worker, args=arg) for arg in worker_args]
[w.start() for w in workers]
[w.join() for w in workers]
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--n_gpus', type=int, default=4)
parser.add_argument('--max_per_gpu', type=int, default=1)
parser.add_argument('--start', type=int, default=0)
parser.add_argument('--end', type=int, default=-1)
args = parser.parse_args()
exps = [
dict(trans_type=['reparam_w_tanh'], batch_size=[128],
lr=[1e-3], weight_decay=[1e-6], z_dim=[8]),
dict(seed=[0, 1, 2, 3]),
dict(root=['data/rope']),
]
exps = construct_variants(exps)
if args.end == -1:
args.end = len(exps)
print(f'Running {args.end - args.start} / {len(exps)} experiments')
exps = exps[args.start:args.end]
n_exps = len(exps)
chunk_size = math.ceil(n_exps / args.n_gpus)
worker_args = []
for i in range(args.n_gpus):
start, end = chunk_size * i, min(chunk_size * (i + 1), n_exps)
worker_args.append((i, args.max_per_gpu, exps[start:end]))
workers = [mp.Process(target=worker, args=arg) for arg in worker_args]
[w.start() for w in workers]
[w.join() for w in workers]