-
Notifications
You must be signed in to change notification settings - Fork 0
/
read_rlds.py
88 lines (73 loc) · 3.1 KB
/
read_rlds.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
import tensorflow_datasets as tfds
import tqdm
import numpy as np
import cv2
import time
import argparse
import matplotlib.pyplot as plt
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--rlds_dir", type=str, default="test_log2")
parser.add_argument("--show_img", action="store_true")
parser.add_argument("--show_figures", action="store_true", help="show images in matplotlib")
parser.add_argument("--replay", action="store_true", help="replay the trajectory on gym env")
parser.add_argument("--ip", type=str, default="ip if we are replaying")
args = parser.parse_args()
ds_builder = tfds.builder_from_directory(args.rlds_dir)
dataset = ds_builder.as_dataset(split='all')
if args.replay:
from manipulator_gym.manipulator_env import ManipulatorEnv
from manipulator_gym.interfaces.interface_service import ActionClientInterface
env = ManipulatorEnv(manipulator_interface=ActionClientInterface(host=args.ip))
# print len of dataset
# print("size of dataset", len(list(dataset)))
# assert len(list(dataset)) == num_of_episodes, f"There should be 3 episodes in the dataset"
# dataset = dataset.repeat().shuffle(1).batch(1)
ds_length = len(list(dataset))
dataset = dataset.take(ds_length)
it = iter(dataset)
for i in tqdm.tqdm(range(ds_length)):
episode = next(it)
print("episode: ", i)
steps = episode['steps']
# for step in steps:
# img = step['observation']['image_primary']
# img = np.array(img)
# cv2.imshow("img", img)
# cv2.waitKey(10)
# break
# continue
print("key in a traj: ", episode.keys())
prim_img_buffer = []
wrist_img_buffer = []
if args.replay:
env.reset()
for step in steps:
print(step['observation'].keys())
# print(step["action"])
# print(" - state: ", step['observation']['state'])
if args.show_img:
img = step['observation']['image_primary']
img = np.array(img)
cv2.imshow("img", img)
cv2.waitKey(10)
if args.show_figures:
prim_img_buffer.append(step['observation']['image_primary'])
wrist_img_buffer.append(step['observation']['image_wrist'])
if args.show_figures and len(prim_img_buffer) == 10:
# show both images in matplot lib with 2 rows and 10 columns
fig, axs = plt.subplots(2, 10)
for i in range(10):
axs[0, i].imshow(prim_img_buffer[i])
axs[1, i].imshow(wrist_img_buffer[i])
plt.show()
prim_img_buffer = []
wrist_img_buffer = []
if args.replay:
action = step['action']
print("replaying action: ", action)
done = env.step(action)
time.sleep(0.1)
print("done")
del it, dataset
cv2.destroyAllWindows()