DB/local_dm_control_suite/utils/parse_amc.py

# 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))