-
Notifications
You must be signed in to change notification settings - Fork 0
/
convert.py
76 lines (63 loc) · 2.19 KB
/
convert.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
import os
import os.path as osp
import numpy as np
import h5py
def save_dict_h5py(data, fname):
"""Save dictionary containing numpy arrays to h5py file."""
with h5py.File(fname, 'w') as hf:
for key in data.keys():
hf.create_dataset(key, data=data[key])
def load_dict_h5py(fname):
"""Restore dictionary containing numpy arrays from h5py file."""
data = dict()
with h5py.File(fname, 'r') as hf:
for key in hf.keys():
data[key] = hf[key][()]
return data
# >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> Ugly, hard-coded conversion to h5.
stems = ['InvertedPendulum-v2',
'InvertedDoublePendulum-v2',
'Reacher-v2',
'Hopper-v3',
'Walker2d-v3',
'HalfCheetah-v3',
'Ant-v3',
'Humanoid-v3']
for stem in stems:
fname = stem + '.trajs_16.npz'
demo_dir = osp.join(os.environ['HOME'],
'Downloads/expert-demos-full-archive/more_demos_npz',
stem)
expert_path = osp.join(demo_dir, fname)
with np.load(expert_path, allow_pickle=True) as data:
data_map = {}
for k, v in data.items():
data_map[k] = v
print(v[0].shape)
new_demo_dir = osp.join('/Users/lionelblonde/Downloads/new_demos_h5', stem)
new_ext = '.h5'
for i in range(16):
data_map_ = {}
for k, v in data_map.items():
if k == 'obs':
k = 'obs0'
elif k == 'next_obs':
k = 'obs1'
elif k == 'dones':
k = 'dones1'
elif k == 'pix_obs':
k = 'pix_obs0'
elif k == 'pix_next_obs':
k = 'pix_obs1'
data_map_.update({k: v[i]})
print(k, v[i].shape)
new_fname = stem + '_demo{}'.format(str(i).zfill(3))
fname = new_fname + new_ext
new_expert_path = osp.join(new_demo_dir, fname)
print(new_expert_path)
os.makedirs(new_demo_dir, exist_ok=True)
save_dict_h5py(data_map_, new_expert_path)
retrieved_data_map = load_dict_h5py(new_expert_path)
for k, v in retrieved_data_map.items():
print(k, v.shape)
print("Bye.")