-
Notifications
You must be signed in to change notification settings - Fork 1
/
mocap_preprocess.py
411 lines (376 loc) · 14.8 KB
/
mocap_preprocess.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
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
"""Preprocessing for embedding motion capture/dannce data."""
import dm_control
import h5py
from dm_control.locomotion.walkers import rodent
from dm_control.locomotion.walkers import rescale
from dm_control.utils import transformations as tr
from dm_control import mjcf
import pickle
import mocap_preprocess
import numpy as np
import sys
import os
import argparse
from scipy.io import loadmat
from typing import Text, List, Tuple, Dict, Union, Optional, Sequence
import subprocess
import jax
from jax import numpy as jp
from flax import struct
import walker
from walker import Rat
from typing import Any
def process(
stac_path: Text,
save_file: Text,
scale_factor: float = 0.9,
start_step: int = 0,
clip_length: int = 250,
n_steps: int = None,
max_qvel: float = 20.0,
dt: float = 0.02,
adjust_z_offset: float = 0.0,
verbatim: bool = False,
ref_steps: Tuple = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
):
"""Process a set of joint angles into the features that
the referenced trajectory is composed of
Args:
stac_path (Text): Path to stac file containing reference.
save_file (Text): Path to Folder in which to save hdf5 dataset.
start_step (int, optional): First frame in rollout
clip_length (int, optional): Length of clip
max_qvel (float, optional): Max allowed qvelocity
dt (float, optional): Timestep
adjust_z_offset (float, optional): Z-offset in m
verbatim (bool, optional): Process in verbatim mode.
ref_steps (Tuple, optional): Reference steps.
"""
# load qpos from file
with open(stac_path, "rb") as file:
d = pickle.load(file)
mocap_qpos = np.array(d["qpos"])
# load rodent mjcf
root = mjcf.from_path("./assets/rodent.xml")
rescale.rescale_subtree(
root,
scale_factor,
scale_factor,
)
mj_model = mjcf.Physics.from_mjcf_model(root).model.ptr
"""Extract featires from the reference qpos"""
if n_steps is None:
n_steps = mocap_qpos.shape[0]
jax_paths = []
max_reference_index = np.max(ref_steps) + 1
with h5py.File(save_file, "w") as file:
for start_step in range(start_step, start_step + n_steps, clip_length):
print(f"start_step: {start_step}", flush=True)
end_step = np.min(
[start_step + clip_length + max_reference_index, start_step + n_steps]
)
mocap_features = get_mocap_features(
mocap_qpos[start_step:end_step, :],
walker,
physics,
max_qvel,
dt,
adjust_z_offset,
verbatim,
)
mocap_features["scaling"] = np.array([])
mocap_features["markers"] = np.array([])
save_features(file, mocap_features, f"clip_{start_step}")
jax_paths.append(
save_dataclass_pickle(
f"{save_file[:-3]}_clip_{start_step}.p", mocap_features
)
)
return jax_paths
def get_mocap_features(
mocap_qpos: np.ndarray,
walker: rodent.Rat,
physics,
max_qvel: float,
dt: float,
adjust_z_offset: float,
verbatim: bool,
null_xyr: bool = False,
shift_position=None,
shift_rotation=None,
):
"""Convert mocap_qpos to valid reference features.
Args:
mocap_qpos (np.ndarray): Array of qpos data
walker (rodent.Rat): rodent walker
physics (TYPE): Environment Physics instance.
max_qvel (float): Maximum allowable q velocity.
dt (float): Timestep between qpos frames.
adjust_z_offset (float): Adjust Z position by this amount.
verbatim (bool): If true, preprocess verbatim.
null_xyr (bool, optional): Description
shift_position (bool, optional): Amount by which to shift position.
shift_rotation (bool, optional): Amount by which to shift the rotation.
"""
# Clip the angles.
joint_names = [b.name for b in walker.mocap_joints]
joint_ranges = physics.bind(walker.mocap_joints).range
min_angles = joint_ranges[:, 0]
max_angles = joint_ranges[:, 1]
angles = mocap_qpos[:, 7:]
clipped_angles = np.clip(angles, min_angles, max_angles)
indexes = np.where(angles != clipped_angles)
if verbatim and indexes[0].size != 0:
for i, j in zip(*indexes):
if np.abs(angles[i, j] - clipped_angles[i, j]) >= 0.1:
print(
"Step {} angle of {} clipped from {} to {}.".format(
i, joint_names[j], angles[i, j], clipped_angles[i, j]
)
)
mocap_qpos[:, 7:] = clipped_angles
# Generate the mocap_features.
mocap_features = {}
mocap_features["position"] = []
mocap_features["quaternion"] = []
mocap_features["joints"] = []
mocap_features["center_of_mass"] = []
mocap_features["end_effectors"] = []
mocap_features["velocity"] = []
mocap_features["angular_velocity"] = []
mocap_features["joints_velocity"] = []
mocap_features["appendages"] = []
mocap_features["body_positions"] = []
mocap_features["body_quaternions"] = []
feet_height = []
walker_bodies = walker.mocap_tracking_bodies
body_names = [b.name for b in walker_bodies]
# print(len(walker_bodies), body_names)
if adjust_z_offset:
left_foot_index = body_names.index("foot_L")
right_foot_index = body_names.index("foot_R")
# Padding for velocity corner case.
mocap_qpos = np.concatenate([mocap_qpos, mocap_qpos[-1, np.newaxis, :]], axis=0)
# print(mocap_qpos.shape)
qvel = np.zeros(len(mocap_qpos[0]) - 1)
for n_frame, qpos in enumerate(mocap_qpos[:-1]):
set_walker(
physics,
walker,
qpos,
qvel,
null_xyr=null_xyr,
position_shift=shift_position,
rotation_shift=shift_rotation,
)
freejoint = walker.mjcf_model.find(
"joint", "root"
) # mjcf.get_attachment_frame(walker.mjcf_model).freejoint
root_pos = physics.bind(freejoint).qpos[:3].copy()
mocap_features["position"].append(root_pos)
root_quat = physics.bind(freejoint).qpos[3:].copy()
mocap_features["quaternion"].append(root_quat)
joints = np.array(physics.bind(walker.mocap_joints).qpos)
mocap_features["joints"].append(joints)
freejoint_frame = walker.mjcf_model.find(
"body", "torso"
) # mjcf.get_attachment_frame(walker.mjcf_model)
com = np.array(physics.bind(freejoint_frame).subtree_com)
mocap_features["center_of_mass"].append(com)
end_effectors = np.copy(
walker.observables.end_effectors_pos(physics)[:]
).reshape(-1, 3)
mocap_features["end_effectors"].append(end_effectors)
if hasattr(walker.observables, "appendages_pos"):
appendages = np.copy(walker.observables.appendages_pos(physics)[:]).reshape(
-1, 3
)
else:
appendages = np.copy(end_effectors)
mocap_features["appendages"].append(appendages)
xpos = physics.bind(walker_bodies).xpos.copy()
mocap_features["body_positions"].append(xpos)
xquat = physics.bind(walker_bodies).xquat.copy()
mocap_features["body_quaternions"].append(xquat)
if adjust_z_offset:
feet_height += [xpos[left_foot_index][2], xpos[right_foot_index][2]]
# Array
mocap_features["position"] = np.array(mocap_features["position"])
mocap_features["quaternion"] = np.array(mocap_features["quaternion"])
mocap_features["joints"] = np.array(mocap_features["joints"])
mocap_features["center_of_mass"] = np.array(mocap_features["center_of_mass"])
mocap_features["end_effectors"] = np.array(mocap_features["end_effectors"])
mocap_features["appendages"] = np.array(mocap_features["appendages"])
mocap_features["body_positions"] = np.array(mocap_features["body_positions"])
mocap_features["body_quaternions"] = np.array(mocap_features["body_quaternions"])
print(mocap_features["position"].shape)
# Offset vertically the qpos and xpos to ensure that the clip is aligned
# with the floor. The heuristic uses the 10 lowest feet heights and
# compensates for the thickness of the geoms.
feet_height = np.sort(feet_height)
if adjust_z_offset:
z_offset = feet_height[:10].mean() - 0.006
else:
z_offset = 0
mocap_qpos[:, 2] -= z_offset
mocap_features["position"][:, 2] -= z_offset
mocap_features["center_of_mass"][:, 2] -= z_offset
mocap_features["body_positions"][:, :, 2] -= z_offset
# Calculate qvel, clip.
mocap_qvel = compute_velocity_from_kinematics(mocap_qpos, dt)
vels = mocap_qvel[:, 6:]
clipped_vels = np.clip(vels, -max_qvel, max_qvel)
indexes = np.where(vels != clipped_vels)
if verbatim and indexes[0].size != 0:
for i, j in zip(*indexes):
if np.abs(vels[i, j] - clipped_vels[i, j]) >= 0.1:
print(
"Step {} velocity of {} clipped from {} to {}.".format(
i, joint_names[j], vels[i, j], clipped_vels[i, j]
)
)
mocap_qvel[:, 6:] = clipped_vels
mocap_features["velocity"] = mocap_qvel[:, :3]
mocap_features["angular_velocity"] = mocap_qvel[:, 3:6]
mocap_features["joints_velocity"] = mocap_qvel[:, 6:]
return mocap_features
def set_walker(
physics,
walker: rodent.Rat,
qpos: np.ndarray,
qvel: np.ndarray,
offset: Union[float, List, np.ndarray] = 0.0,
null_xyr: bool = False,
position_shift=None,
rotation_shift=None,
):
"""Set the freejoint and walker's joints angles and velocities.
Args:
physics (TYPE): Environment Physics instance.
walker (rodent.Rat): Description
qpos (np.ndarray): Description
qvel (np.ndarray): Description
offset (Union[float, List, np.ndarray], optional): xyz offset
null_xyr (bool, optional): Description
position_shift (TYPE, optional): Amount by which to shift position.
rotation_shift (TYPE, optional): Amount by which to shift the rotation.
"""
qpos = qpos.copy()
if null_xyr:
qpos[:3] = 0.0
euler = tr.quat_to_euler(qpos[3:7], ordering="ZYX")
euler[0] = 0.0
quat = tr.euler_to_quat(euler, ordering="ZYX")
qpos[3:7] = quat
qpos[:3] += offset
freejoint = walker.mjcf_model.find(
"joint", "root"
) # mjcf.get_attachment_frame(walker.mjcf_model).freejoint
physics.bind(freejoint).qpos = qpos[:7]
physics.bind(freejoint).qvel = qvel[:6]
physics.bind(walker.mocap_joints).qpos = qpos[7:]
physics.bind(walker.mocap_joints).qvel = qvel[6:]
if position_shift is not None or rotation_shift is not None:
walker.shift_pose(
physics,
position=position_shift,
quaternion=rotation_shift,
rotate_velocity=True,
)
def compute_velocity_from_kinematics(
qpos_trajectory: np.ndarray, dt: float
) -> np.ndarray:
"""Computes velocity trajectory from position trajectory.
Args:
qpos_trajectory (np.ndarray): trajectory of qpos values T x ?
Note assumes has freejoint as the first 7 dimensions
dt (float): timestep between qpos entries
Returns:
np.ndarray: Trajectory of velocities.
"""
qvel_translation = (qpos_trajectory[1:, :3] - qpos_trajectory[:-1, :3]) / dt
qvel_gyro = []
for t in range(qpos_trajectory.shape[0] - 1):
normed_diff = tr.quat_diff(qpos_trajectory[t, 3:7], qpos_trajectory[t + 1, 3:7])
normed_diff /= np.linalg.norm(normed_diff)
qvel_gyro.append(tr.quat_to_axisangle(normed_diff) / dt)
qvel_gyro = np.stack(qvel_gyro)
qvel_joints = (qpos_trajectory[1:, 7:] - qpos_trajectory[:-1, 7:]) / dt
return np.concatenate([qvel_translation, qvel_gyro, qvel_joints], axis=1)
# 13 features
@struct.dataclass
class ReferenceClip:
angular_velocity: jp.ndarray
appendages: jp.ndarray
body_positions: jp.ndarray
body_quaternions: jp.ndarray
center_of_mass: jp.ndarray
end_effectors: jp.ndarray
joints: jp.ndarray
joints_velocity: jp.ndarray
markers: jp.ndarray
position: jp.ndarray
quaternion: jp.ndarray
scaling: jp.ndarray
velocity: jp.ndarray
class ClipCollection:
"""Dataclass representing a collection of mocap reference clips."""
def __init__(
self,
ids: Sequence[Text],
start_steps: Optional[Sequence[int]] = None,
end_steps: Optional[Sequence[int]] = None,
weights: Optional[Sequence[Union[int, float]]] = None,
):
"""Instantiate a ClipCollection."""
self.ids = ids
self.start_steps = start_steps
self.end_steps = end_steps
self.weights = weights
num_clips = len(self.ids)
try:
if self.start_steps is None:
# by default start at the beginning
self.start_steps = (0,) * num_clips
else:
assert len(self.start_steps) == num_clips
# without access to the actual clip we cannot specify an end_steps default
if self.end_steps is not None:
assert len(self.end_steps) == num_clips
if self.weights is None:
self.weights = (1.0,) * num_clips
else:
assert len(self.weights) == num_clips
assert jp.all(np.array(self.weights) >= 0.0)
except AssertionError as e:
raise ValueError("ClipCollection validation failed. {}".format(e))
def save_dataclass_pickle(pickle_path, mocap_features):
data = ReferenceClip(**mocap_features)
data = jax.tree_map(lambda x: jp.array(x), data)
with open(pickle_path, "wb") as f:
pickle.dump(data, f)
return pickle_path
def save_features(file: h5py.File, mocap_features: Dict, clip_name: Text):
"""Save features to hdf5 dataset
Args:
file (h5py.File): Hdf5 dataset
mocap_features (Dict): Features extracted through rollout
clip_name (Text): Name of the clip stored in the hdf5 dataset.
"""
clip_group = file.create_group(clip_name)
n_steps = len(mocap_features["center_of_mass"])
clip_group.attrs["num_steps"] = n_steps
clip_group.attrs["dt"] = 0.02
file.create_group("/" + clip_name + "/walkers")
file.create_group("/" + clip_name + "/props")
walker_group = file.create_group("/" + clip_name + "/walkers/walker_0")
for k, v in mocap_features.items():
if len(np.array(v).shape) == 3:
v = np.transpose(v, (1, 2, 0))
# print(v.shape)
walker_group[k] = np.reshape(np.array(v), (-1, n_steps))
elif len(np.array(v).shape) == 2:
v = np.swapaxes(v, 0, 1)
walker_group[k] = v
else:
walker_group[k] = v