-
Notifications
You must be signed in to change notification settings - Fork 10
/
generate_data.py
177 lines (154 loc) · 7.17 KB
/
generate_data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
import argparse
import ast
import copy
import shutil
import sys
import os
import yaml
import numpy as np
from tqdm import tqdm
from environments import environment_factory
def generate_and_save(root_path, environment, n_samples, n_frames, delta_time, img_size,
radius_bound, noise_level, color, start_seed, train=True):
path = os.path.join(root_path, 'train' if train else 'test')
if not os.path.exists(path):
os.makedirs(path)
for i in tqdm(range(n_samples)):
rolls = environment.sample_random_rollouts(
number_of_frames=n_frames,
delta_time=delta_time,
number_of_rollouts=1,
img_size=img_size,
noise_level=noise_level,
radius_bound=radius_bound,
color=color,
seed=i + start_seed
)[0]
filename = "{0:05d}".format(i)
np.savez(os.path.join(path, filename), rolls)
return path
def _read_config(config_file):
with open(config_file, 'r') as file:
config = yaml.load(file, Loader=yaml.FullLoader)
return config
def _prepare_out_config(config, train_path, test_path):
out_config = copy.deepcopy(config)
out_config['dataset']['train_data'] = train_path
out_config['dataset']['test_data'] = test_path
return out_config
def _overwrite_config_with_cmd_arguments(config, args):
# This function overwrites parameters in the given dictionary
# with the correspondent command line arguments.
if args.ntrain is not None:
config['dataset']['num_train_samples'] = args.ntrain[0]
if args.ntest is not None:
config['dataset']['num_test_samples'] = args.ntest[0]
if args.env is not None:
env_params = _read_config(DEFAULT_ENVIRONMENTS_PATH + args.env[0] + '.yaml')
config['environment'] = env_params['environment']
if args.params is not None:
for p in args.params:
key, value = p.split('=')
ptr = config
keys = key.split('.')
for i, k in enumerate(keys):
if i == len(keys) - 1:
ptr[k] = ast.literal_eval(value)
else:
ptr = ptr[k]
if __name__ == '__main__':
DEFAULT_DATASETS_ROOT = 'datasets/'
DEFAULT_DATASET_CONFIG_FILE = 'experiment_params/dataset_online_default.yaml'
DEFAULT_TRAIN_CONFIG_FILE = 'experiment_params/train_config_default.yaml'
DEFAULT_ENVIRONMENTS_PATH = 'experiment_params/default_environments/'
parser = argparse.ArgumentParser()
parser.add_argument(
'--name', action='store', nargs=1, required=True, help='The dataset name.'
)
parser.add_argument(
'--dataset-config', action='store', nargs=1, type=str, required=False,
help=f'YAML file from which to read the dataset parameters. If not specified, '
f'{DEFAULT_DATASET_CONFIG_FILE} will be used.'
)
parser.add_argument(
'--ntrain', action='store', nargs=1, required=False, type=int,
help='Number of training sample to generate.'
)
parser.add_argument(
'--ntest', action='store', nargs=1, required=False, type=int,
help='Number of test samples to generate.'
)
parser.add_argument(
'--env', action='store', nargs=1, required=False, type=str,
help=f'The default environment specifications to use. Can be \'pendulum\', \'spring\', '
f'\'two_bodies\', \'three_bodies\', \'chaotic_pendulum\'. If this argument is '
f'specified, a default environment section will be loaded from the correspondent yaml '
f'file in {DEFAULT_ENVIRONMENTS_PATH}'
)
parser.add_argument(
'--datasets-root', action='store', nargs=1, required=False, type=str,
help=f'Root of the datasets folder in which the dataset will be stored. If not specified, '
f'{DEFAULT_DATASETS_ROOT} will be used as default.'
)
parser.add_argument(
'--params', action='store', nargs='+', required=False,
help='Override one or more parameters in the config. The format of an argument is '
'param_name=param_value. Nested parameters are accessible by using a dot, '
'i.e. --param dataset.img_size=32. IMPORTANT: lists must be enclosed in double '
'quotes, i.e. --param environment.mass:"[0.5, 0.5]".'
)
_args = parser.parse_args()
# Read yaml file with parameters definition
_dataset_config_file = _args.dataset_config[0] if _args.dataset_config is not None else \
DEFAULT_DATASET_CONFIG_FILE
_dataset_config = _read_config(_dataset_config_file)
# Overwrite dictionary from command line args to ensure they will be used
_overwrite_config_with_cmd_arguments(_dataset_config, _args)
# Extract environment parameters
EXP_NAME = _args.name[0]
N_TRAIN_SAMPLES = _dataset_config['dataset']['num_train_samples']
N_TEST_SAMPLES = _dataset_config['dataset']['num_test_samples']
IMG_SIZE = _dataset_config['dataset']['img_size']
RADIUS_BOUND = _dataset_config['dataset']['radius_bound']
NOISE_LEVEL = _dataset_config['dataset']['rollout']['noise_level']
N_FRAMES = _dataset_config['dataset']['rollout']['seq_length']
DELTA_TIME = _dataset_config['dataset']['rollout']['delta_time']
N_CHANNELS = _dataset_config['dataset']['rollout']['n_channels']
# Get dataset output path
dataset_root = DEFAULT_DATASETS_ROOT if _args.datasets_root is None else _args.datasets_root[0]
dataset_root = os.path.join(dataset_root, EXP_NAME)
# Get the environment object from dictionary parameters
environment = environment_factory.EnvFactory.get_environment(**_dataset_config['environment'])
# Ask user confirmation
print(f'The dataset will be generated with the following configuration:')
print(f'PATH: {dataset_root}')
print(f'dataset: {_dataset_config["dataset"]}')
print(f'environment: {_dataset_config["environment"]}')
print('\nProceed? (y/n):')
if input() != 'y':
print('Aborting')
exit()
# Generate train samples
_train_path = generate_and_save(
root_path=dataset_root, environment=environment,
n_samples=N_TRAIN_SAMPLES, n_frames=N_FRAMES, delta_time=DELTA_TIME, img_size=IMG_SIZE,
radius_bound=RADIUS_BOUND, noise_level=NOISE_LEVEL, color=N_CHANNELS == 3,
start_seed=0, train=True
)
# Generate test samples
_test_path = None
if N_TEST_SAMPLES > 0:
_test_path = generate_and_save(
root_path=dataset_root, environment=environment,
n_samples=N_TEST_SAMPLES, n_frames=N_FRAMES, delta_time=DELTA_TIME,
img_size=IMG_SIZE, radius_bound=RADIUS_BOUND, noise_level=NOISE_LEVEL,
color=N_CHANNELS == 3, start_seed=N_TRAIN_SAMPLES, train=False
)
# Convert parameters to offline train parameters and write them in the dataset
_out_config = _prepare_out_config(_dataset_config, _train_path, _test_path)
yaml_content = yaml.dump(_out_config, default_flow_style=True)
config_out_path = os.path.join(dataset_root, 'parameters.yaml')
with open(config_out_path, 'x') as f:
f.write(yaml_content)
f.close()
print(f'A parameter file ready to be trained on was generated at {config_out_path}')