72 lines
1.7 KiB
Python
Raw Permalink Normal View History

2025-07-28 12:35:44 +08:00
import argparse
import pickle
import matplotlib.pyplot as plt
import numpy as np
parser = argparse.ArgumentParser()
parser.add_argument("-d", "--data", type=str, required=True)
args = parser.parse_args()
data = pickle.load(open(args.data, "rb"))
if "robot" in data and "mujoco" in data:
res = None
while res is None or res not in ["1", "2"]:
res = input("plot robot (1) or mujoco (2) ? ")
if res == "1":
command_value = data["robot"]
title = "Robot"
else:
command_value = data["mujoco"]
title = "Mujoco"
elif "mujoco" in data:
command_value = data["mujoco"]
title = "Mujoco"
elif "robot" in data:
command_value = data["robot"]
title = "Robot"
else:
print("NO DATA")
exit()
dofs = {
0: "right_hip_yaw",
1: "right_hip_roll",
2: "right_hip_pitch",
3: "right_knee",
4: "right_ankle",
5: "left_hip_yaw",
6: "left_hip_roll",
7: "left_hip_pitch",
8: "left_knee",
9: "left_ankle",
10: "neck_pitch",
11: "head_pitch",
12: "head_yaw",
13: "left_antenna",
14: "right_antenna",
}
# command_value = np.array(command_value)
fig, axs = plt.subplots(4, 4)
dof_id = 0
for i in range(4):
for j in range(4):
if i == 3 and j == 3:
break
print(4 * i + j)
command = []
value = []
for k in range(len(command_value)):
command.append(command_value[k][0][4 * i + j])
value.append(command_value[k][1][4 * i + j])
axs[i, j].plot(command, label="command")
axs[i, j].plot(value, label="value")
axs[i, j].legend()
axs[i, j].set_title(f"{dofs[dof_id]}")
dof_id += 1
name = args.data.split("/")[-1].split(".")[0]
fig.suptitle(title)
plt.show()