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]