def convert()

in local_dm_control_suite/utils/parse_amc.py [0:0]

• 37 lines of code
• 9 McCabe index (conditional complexity)

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)