-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathdataset_paths.py
More file actions
17 lines (15 loc) · 882 Bytes
/
Copy pathdataset_paths.py
File metadata and controls
17 lines (15 loc) · 882 Bytes
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import os
##### paths to datasets for training diffusor models
def get_diffuser_dataset_path(dataset_loadbase, num_entity, input_type, rand_color=False, push_t=False, kitchen=False):
if input_type == 'dlp':
if kitchen:
assert num_entity == 1
return os.path.join(dataset_loadbase, 'kitchen', 'dlp_kitchen_dataset_40kp_64kpp_4zdim.pkl')
if push_t:
return os.path.join(dataset_loadbase, 'push_t', f'{num_entity}T', 'panda_push_replay_buffer_dlp.pkl')
else:
if rand_color:
return os.path.join(dataset_loadbase, 'push_cubes', f'{num_entity}C_randcolor', 'panda_push_replay_buffer_dlp.pkl')
else:
return os.path.join(dataset_loadbase, 'push_cubes', f'{num_entity}C', 'panda_push_replay_buffer_dlp.pkl')
raise ValueError(f'input_type {input_type} not supported')