252 lines
8.0 KiB
Python
252 lines
8.0 KiB
Python
|
# Copyright 2017 The dm_control Authors.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
# ============================================================================
|
||
|
|
|
||
|
|
"""Parse and convert amc motion capture data."""
|
||
|
|
|
||
|
|
from __future__ import absolute_import
|
||
|
|
from __future__ import division
|
||
|
|
from __future__ import print_function
|
||
|
|
|
||
|
|
import collections
|
||
|
|
|
||
|
|
from dm_control.mujoco.wrapper import mjbindings
|
||
|
|
import numpy as np
|
||
|
|
from scipy import interpolate
|
||
|
|
from six.moves import range
|
||
|
|
|
||
|
|
mjlib = mjbindings.mjlib
|
||
|
|
|
||
|
|
MOCAP_DT = 1.0/120.0
|
||
|
|
CONVERSION_LENGTH = 0.056444
|
||
|
|
|
||
|
|
_CMU_MOCAP_JOINT_ORDER = (
|
||
|
|
'root0', 'root1', 'root2', 'root3', 'root4', 'root5', 'lowerbackrx',
|
||
|
|
'lowerbackry', 'lowerbackrz', 'upperbackrx', 'upperbackry', 'upperbackrz',
|
||
|
|
'thoraxrx', 'thoraxry', 'thoraxrz', 'lowerneckrx', 'lowerneckry',
|
||
|
|
'lowerneckrz', 'upperneckrx', 'upperneckry', 'upperneckrz', 'headrx',
|
||
|
|
'headry', 'headrz', 'rclaviclery', 'rclaviclerz', 'rhumerusrx',
|
||
|
|
'rhumerusry', 'rhumerusrz', 'rradiusrx', 'rwristry', 'rhandrx', 'rhandrz',
|
||
|
|
'rfingersrx', 'rthumbrx', 'rthumbrz', 'lclaviclery', 'lclaviclerz',
|
||
|
|
'lhumerusrx', 'lhumerusry', 'lhumerusrz', 'lradiusrx', 'lwristry',
|
||
|
|
'lhandrx', 'lhandrz', 'lfingersrx', 'lthumbrx', 'lthumbrz', 'rfemurrx',
|
||
|
|
'rfemurry', 'rfemurrz', 'rtibiarx', 'rfootrx', 'rfootrz', 'rtoesrx',
|
||
|
|
'lfemurrx', 'lfemurry', 'lfemurrz', 'ltibiarx', 'lfootrx', 'lfootrz',
|
||
|
|
'ltoesrx'
|
||
|
|
)
|
||
|
|
|
||
|
|
Converted = collections.namedtuple('Converted',
|
||
|
|
['qpos', 'qvel', 'time'])
|
||
|
|
|
||
|
|
|
||
|
|
def convert(file_name, physics, timestep):
|
||
|
|
"""Converts the parsed .amc values into qpos and qvel values and resamples.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
file_name: The .amc file to be parsed and converted.
|
||
|
|
physics: The corresponding physics instance.
|
||
|
|
timestep: Desired output interval between resampled frames.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A namedtuple with fields:
|
||
|
|
`qpos`, a numpy array containing converted positional variables.
|
||
|
|
`qvel`, a numpy array containing converted velocity variables.
|
||
|
|
`time`, a numpy array containing the corresponding times.
|
||
|
|
"""
|
||
|
|
frame_values = parse(file_name)
|
||
|
|
joint2index = {}
|
||
|
|
for name in physics.named.data.qpos.axes.row.names:
|
||
|
|
joint2index[name] = physics.named.data.qpos.axes.row.convert_key_item(name)
|
||
|
|
index2joint = {}
|
||
|
|
for joint, index in joint2index.items():
|
||
|
|
if isinstance(index, slice):
|
||
|
|
indices = range(index.start, index.stop)
|
||
|
|
else:
|
||
|
|
indices = [index]
|
||
|
|
for ii in indices:
|
||
|
|
index2joint[ii] = joint
|
||
|
|
|
||
|
|
# Convert frame_values to qpos
|
||
|
|
amcvals2qpos_transformer = Amcvals2qpos(index2joint, _CMU_MOCAP_JOINT_ORDER)
|
||
|
|
qpos_values = []
|
||
|
|
for frame_value in frame_values:
|
||
|
|
qpos_values.append(amcvals2qpos_transformer(frame_value))
|
||
|
|
qpos_values = np.stack(qpos_values) # Time by nq
|
||
|
|
|
||
|
|
# Interpolate/resample.
|
||
|
|
# Note: interpolate quaternions rather than euler angles (slerp).
|
||
|
|
# see https://en.wikipedia.org/wiki/Slerp
|
||
|
|
qpos_values_resampled = []
|
||
|
|
time_vals = np.arange(0, len(frame_values)*MOCAP_DT - 1e-8, MOCAP_DT)
|
||
|
|
time_vals_new = np.arange(0, len(frame_values)*MOCAP_DT, timestep)
|
||
|
|
while time_vals_new[-1] > time_vals[-1]:
|
||
|
|
time_vals_new = time_vals_new[:-1]
|
||
|
|
|
||
|
|
for i in range(qpos_values.shape[1]):
|
||
|
|
f = interpolate.splrep(time_vals, qpos_values[:, i])
|
||
|
|
qpos_values_resampled.append(interpolate.splev(time_vals_new, f))
|
||
|
|
|
||
|
|
qpos_values_resampled = np.stack(qpos_values_resampled) # nq by ntime
|
||
|
|
|
||
|
|
qvel_list = []
|
||
|
|
for t in range(qpos_values_resampled.shape[1]-1):
|
||
|
|
p_tp1 = qpos_values_resampled[:, t + 1]
|
||
|
|
p_t = qpos_values_resampled[:, t]
|
||
|
|
qvel = [(p_tp1[:3]-p_t[:3])/ timestep,
|
||
|
|
mj_quat2vel(mj_quatdiff(p_t[3:7], p_tp1[3:7]), timestep),
|
||
|
|
(p_tp1[7:]-p_t[7:])/ timestep]
|
||
|
|
qvel_list.append(np.concatenate(qvel))
|
||
|
|
|
||
|
|
qvel_values_resampled = np.vstack(qvel_list).T
|
||
|
|
|
||
|
|
return Converted(qpos_values_resampled, qvel_values_resampled, time_vals_new)
|
||
|
|
|
||
|
|
|
||
|
|
def parse(file_name):
|
||
|
|
"""Parses the amc file format."""
|
||
|
|
values = []
|
||
|
|
fid = open(file_name, 'r')
|
||
|
|
line = fid.readline().strip()
|
||
|
|
frame_ind = 1
|
||
|
|
first_frame = True
|
||
|
|
while True:
|
||
|
|
# Parse first frame.
|
||
|
|
if first_frame and line[0] == str(frame_ind):
|
||
|
|
first_frame = False
|
||
|
|
frame_ind += 1
|
||
|
|
frame_vals = []
|
||
|
|
while True:
|
||
|
|
line = fid.readline().strip()
|
||
|
|
if not line or line == str(frame_ind):
|
||
|
|
values.append(np.array(frame_vals, dtype=np.float))
|
||
|
|
break
|
||
|
|
tokens = line.split()
|
||
|
|
frame_vals.extend(tokens[1:])
|
||
|
|
# Parse other frames.
|
||
|
|
elif line == str(frame_ind):
|
||
|
|
frame_ind += 1
|
||
|
|
frame_vals = []
|
||
|
|
while True:
|
||
|
|
line = fid.readline().strip()
|
||
|
|
if not line or line == str(frame_ind):
|
||
|
|
values.append(np.array(frame_vals, dtype=np.float))
|
||
|
|
break
|
||
|
|
tokens = line.split()
|
||
|
|
frame_vals.extend(tokens[1:])
|
||
|
|
else:
|
||
|
|
line = fid.readline().strip()
|
||
|
|
if not line:
|
||
|
|
break
|
||
|
|
return values
|
||
|
|
|
||
|
|
|
||
|
|
class Amcvals2qpos(object):
|
||
|
|
"""Callable that converts .amc values for a frame and to MuJoCo qpos format.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, index2joint, joint_order):
|
||
|
|
"""Initializes a new Amcvals2qpos instance.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
index2joint: List of joint angles in .amc file.
|
||
|
|
joint_order: List of joint names in MuJoco MJCF.
|
||
|
|
"""
|
||
|
|
# Root is x,y,z, then quat.
|
||
|
|
# need to get indices of qpos that order for amc default order
|
||
|
|
self.qpos_root_xyz_ind = [0, 1, 2]
|
||
|
|
self.root_xyz_ransform = np.array(
|
||
|
|
[[1, 0, 0], [0, 0, -1], [0, 1, 0]]) * CONVERSION_LENGTH
|
||
|
|
self.qpos_root_quat_ind = [3, 4, 5, 6]
|
||
|
|
amc2qpos_transform = np.zeros((len(index2joint), len(joint_order)))
|
||
|
|
for i in range(len(index2joint)):
|
||
|
|
for j in range(len(joint_order)):
|
||
|
|
if index2joint[i] == joint_order[j]:
|
||
|
|
if 'rx' in index2joint[i]:
|
||
|
|
amc2qpos_transform[i][j] = 1
|
||
|
|
elif 'ry' in index2joint[i]:
|
||
|
|
amc2qpos_transform[i][j] = 1
|
||
|
|
elif 'rz' in index2joint[i]:
|
||
|
|
amc2qpos_transform[i][j] = 1
|
||
|
|
self.amc2qpos_transform = amc2qpos_transform
|
||
|
|
|
||
|
|
def __call__(self, amc_val):
|
||
|
|
"""Converts a `.amc` frame to MuJoCo qpos format."""
|
||
|
|
amc_val_rad = np.deg2rad(amc_val)
|
||
|
|
qpos = np.dot(self.amc2qpos_transform, amc_val_rad)
|
||
|
|
|
||
|
|
# Root.
|
||
|
|
qpos[:3] = np.dot(self.root_xyz_ransform, amc_val[:3])
|
||
|
|
qpos_quat = euler2quat(amc_val[3], amc_val[4], amc_val[5])
|
||
|
|
qpos_quat = mj_quatprod(euler2quat(90, 0, 0), qpos_quat)
|
||
|
|
|
||
|
|
for i, ind in enumerate(self.qpos_root_quat_ind):
|
||
|
|
qpos[ind] = qpos_quat[i]
|
||
|
|
|
||
|
|
return qpos
|
||
|
|
|
||
|
|
|
||
|
|
def euler2quat(ax, ay, az):
|
||
|
|
"""Converts euler angles to a quaternion.
|
||
|
|
|
||
|
|
Note: rotation order is zyx
|
||
|
|
|
||
|
|
Args:
|
||
|
|
ax: Roll angle (deg)
|
||
|
|
ay: Pitch angle (deg).
|
||
|
|
az: Yaw angle (deg).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
A numpy array representing the rotation as a quaternion.
|
||
|
|
"""
|
||
|
|
r1 = az
|
||
|
|
r2 = ay
|
||
|
|
r3 = ax
|
||
|
|
|
||
|
|
c1 = np.cos(np.deg2rad(r1 / 2))
|
||
|
|
s1 = np.sin(np.deg2rad(r1 / 2))
|
||
|
|
c2 = np.cos(np.deg2rad(r2 / 2))
|
||
|
|
s2 = np.sin(np.deg2rad(r2 / 2))
|
||
|
|
c3 = np.cos(np.deg2rad(r3 / 2))
|
||
|
|
s3 = np.sin(np.deg2rad(r3 / 2))
|
||
|
|
|
||
|
|
q0 = c1 * c2 * c3 + s1 * s2 * s3
|
||
|
|
q1 = c1 * c2 * s3 - s1 * s2 * c3
|
||
|
|
q2 = c1 * s2 * c3 + s1 * c2 * s3
|
||
|
|
q3 = s1 * c2 * c3 - c1 * s2 * s3
|
||
|
|
|
||
|
|
return np.array([q0, q1, q2, q3])
|
||
|
|
|
||
|
|
|
||
|
|
def mj_quatprod(q, r):
|
||
|
|
quaternion = np.zeros(4)
|
||
|
|
mjlib.mju_mulQuat(quaternion, np.ascontiguousarray(q),
|
||
|
|
np.ascontiguousarray(r))
|
||
|
|
return quaternion
|
||
|
|
|
||
|
|
|
||
|
|
def mj_quat2vel(q, dt):
|
||
|
|
vel = np.zeros(3)
|
||
|
|
mjlib.mju_quat2Vel(vel, np.ascontiguousarray(q), dt)
|
||
|
|
return vel
|
||
|
|
|
||
|
|
|
||
|
|
def mj_quatneg(q):
|
||
|
|
quaternion = np.zeros(4)
|
||
|
|
mjlib.mju_negQuat(quaternion, np.ascontiguousarray(q))
|
||
|
|
return quaternion
|
||
|
|
|
||
|
|
|
||
|
|
def mj_quatdiff(source, target):
|
||
|
|
return mj_quatprod(mj_quatneg(source), np.ascontiguousarray(target))
|