-
Notifications
You must be signed in to change notification settings - Fork 8
/
Copy pathsim.py
257 lines (213 loc) · 9.02 KB
/
sim.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
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
import os
import numpy as np
import time
import json
import mujoco
import mujoco.viewer
import placo
import pandas
from placo_utils.tf import tf
from bam.model import load_model
from bam.mujoco import MujocoController
class MujocoSimulation2R:
def __init__(self, testbench: str):
"""
Loading the 2R simulation
"""
this_directory = os.path.dirname(os.path.realpath(__file__))
self.model: mujoco.MjModel = mujoco.MjModel.from_xml_path(
f"{this_directory}/2r_{testbench}/scene.xml"
)
self.data: mujoco.MjData = mujoco.MjData(self.model)
self.testbench = testbench
# Placo robot
self.robot = None
self.viewer = None
self.viewer_start = None
self.t: float = 0
self.dt: float = self.model.opt.timestep
self.frame: int = 0
def step(self) -> None:
self.t = self.frame * self.dt
mujoco.mj_step(self.model, self.data)
self.frame += 1
def reset(self):
self.t = 0
self.frame = 0
self.viewer_start = time.time()
def render(self, realtime: bool = True):
"""
Renders the visualization of the simulation.
Args:
realtime (bool, optional): if True, render will sleep to ensure real time viewing. Defaults to True.
"""
if self.viewer is None:
self.viewer = mujoco.viewer.launch_passive(self.model, self.data)
self.viewer_start = time.time()
if realtime:
current_ts = self.viewer_start + self.frame * self.dt
to_sleep = current_ts - time.time()
if to_sleep > 0:
time.sleep(to_sleep)
self.viewer.sync()
def simulate_log(
self, data: dict, params: str, replay: bool = False, render: bool = False
):
if self.robot is None:
this_directory = os.path.dirname(os.path.realpath(__file__))
self.robot = placo.RobotWrapper(
this_directory + f"/2r_{self.testbench}/robot.urdf", placo.Flags.ignore_collisions
)
if self.testbench in ["mx"]:
self.robot.set_T_world_frame("base", tf.rotation_matrix(np.pi, [1, 0, 0]))
# Updating actuator KP
if "," in params:
params_r1, params_r2 = params.split(",")
model_r1, model_r2 = load_model(params_r1), load_model(params_r2)
else:
model_r1, model_r2 = load_model(params), load_model(params)
if type(data["kp"]) is list:
model_r1.actuator.kp = data["kp"][0]
model_r2.actuator.kp = data["kp"][1]
else:
model_r1.actuator.kp = data["kp"]
model_r2.actuator.kp = data["kp"]
# Creating bam controllers
if not replay:
r1 = MujocoController(model_r1, "R1", self.model, self.data)
r2 = MujocoController(model_r2, "R2", self.model, self.data)
# Setting initial configuration
self.data.joint("R1").qpos[0] = data["entries"][0]["r1"]["position"]
self.data.joint("R2").qpos[0] = data["entries"][0]["r2"]["position"]
log_t0 = data["entries"][0]["timestamp"]
self.reset()
entry_index = 0
running = True
while running:
entry = data["entries"][entry_index]
self.step()
if render:
self.render()
if replay:
# If it's a replay, simply jump to the read position
self.data.joint("R1").qpos[0] = entry["r1"]["position"]
self.data.joint("R2").qpos[0] = entry["r2"]["position"]
else:
r1.update(entry["r1"]["goal_position"])
r2.update(entry["r2"]["goal_position"])
while running and (log_t0 + self.t >= entry["timestamp"]):
entry = data["entries"][entry_index]
entry_index += 1
entry["r1"]["sim_position"] = self.data.joint("R1").qpos[0]
entry["r2"]["sim_position"] = self.data.joint("R2").qpos[0]
entry["end_effector"] = {}
for position in "position", "goal_position", "sim_position":
self.robot.set_joint("R1", entry["r1"][position])
self.robot.set_joint("R2", entry["r2"][position])
self.robot.update_kinematics()
pos = self.robot.get_T_world_frame("end")[:3, 3]
entry["end_effector"][position] = pos
if entry_index == len(data["entries"]):
running = False
if __name__ == "__main__":
import argparse
import matplotlib.pyplot as plt
args_parser = argparse.ArgumentParser()
args_parser.add_argument("--log", type=str, default="2R/log.json", nargs="+")
args_parser.add_argument("--params", type=str, default=[], nargs="+")
args_parser.add_argument("--testbench", type=str, required=True)
args_parser.add_argument("--replay", action="store_true")
args_parser.add_argument("--render", action="store_true")
args_parser.add_argument("--plot", action="store_true")
args_parser.add_argument("--plot_joint", action="store_true")
args_parser.add_argument("--vertical", action="store_true")
args_parser.add_argument("--mae", action="store_true")
args = args_parser.parse_args()
# Loading bam model
sim = MujocoSimulation2R(testbench=args.testbench)
maes = {}
for log in args.log:
# Loading log
data = json.load(open(log))
maes[log] = {}
n = len(args.params)
if args.plot:
# Creating n horizontal subplots
if args.vertical:
f, axs = plt.subplots(n, 1, sharex=True)
else:
f, axs = plt.subplots(1, n, sharey=True)
if n == 1:
axs = [axs]
# Setting figure size
f.set_size_inches(12, 4)
else:
axs = [None] * n
for params, ax in zip(args.params, axs):
sim.simulate_log(data, params, args.replay, args.render)
mae = 0
for dof in "r1", "r2":
errors = [
entry[dof]["position"] - entry[dof]["sim_position"]
for entry in data["entries"]
]
mae += np.mean(np.abs(errors))
mae /= 2
maes[log][params] = mae
if args.plot:
for position in "position", "goal_position", "sim_position":
ax.plot(
[entry["end_effector"][position][0] for entry in data["entries"]],
[entry["end_effector"][position][2] for entry in data["entries"]],
label=position,
ls="--" if position == "goal_position" else "-",
)
ax.legend()
ax.grid()
ax.set_aspect('equal',adjustable='box')
ax.set_title(f"{os.path.basename(log)}, {params}")
if args.plot_joint:
for dof in "r1", "r2":
# Creating two subplots axises
f, (ax1, ax2) = plt.subplots(2, sharex=True)
goal_positions = [
entry[dof]["goal_position"] for entry in data["entries"]
]
positions = [entry[dof]["position"] for entry in data["entries"]]
sim_positions = [
entry[dof]["sim_position"] for entry in data["entries"]
]
ax1.plot(goal_positions, label=f"{dof} goal", color="red")
ax1.plot(positions, label=f"{dof} read", color="blue")
ax1.plot(sim_positions, label=f"{dof} sim", color="green")
ax1.grid()
ax1.legend()
errors = [read - sim for read, sim in zip(positions, sim_positions)]
mae = np.mean(np.abs(errors))
print("MAE: ", mae)
ax2.plot(errors, color="black", label="Simulation error")
ax2.set_ylim(-0.05, 0.05)
ax2.grid()
plt.title(f"{log}, {params}")
plt.show()
if args.plot:
plt.tight_layout()
plt.show()
if args.mae:
total_mae = {params: [] for params in args.params}
for log in maes:
print(f"Log: {log}")
for params in maes[log]:
print(f" {params}: {maes[log][params]}")
total_mae[params] += [maes[log][params]]
labels = [os.path.basename(log) for log in maes]
df = pandas.DataFrame(total_mae, index=labels)
df.plot(kind="bar")
plt.grid(axis="y")
# Setting x label with 45°, keeping the top aligned
plt.xticks(rotation=45, ha="right")
plt.title("MAE per log")
plt.tight_layout()
plt.show()
for params in total_mae:
print(f"Total MAE for {params}: {np.mean(total_mae[params])}")