Added math.py with utility functions for quaternions. Refactored parse_amc.py to depend on math.py.

DeepMind · copybara-github · commit 6604dce21cf0 · 2019-12-20T09:48:49.000-08:00
PiperOrigin-RevId: 286594366
Change-Id: I55e9843680c6e9fdc98c3b5904971eb3603fc9a4
diff --git a/dm_control/mujoco/math.py b/dm_control/mujoco/math.py
@@ -0,0 +1,79 @@
+# Copyright 2019 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.
+# ============================================================================
+
+"""Utility functions."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from dm_control.mujoco.wrapper.mjbindings import mjlib
+
+import numpy as np
+
+
+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))
diff --git a/dm_control/mujoco/math_test.py b/dm_control/mujoco/math_test.py
@@ -0,0 +1,61 @@
+# Copyright 2019 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.
+# ============================================================================
+
+"""Tests for index."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import math
+from absl.testing import absltest
+from absl.testing import parameterized
+from dm_control.mujoco import math as mjmath
+import numpy as np
+
+
+class MathTest(parameterized.TestCase):
+
+  def testQuatProd(self):
+    np.testing.assert_allclose(
+        mjmath.mj_quatprod([0., 1., 0., 0.], [0., 0., 1., 0.]),
+        [0., 0., 0., 1.])
+    np.testing.assert_allclose(
+        mjmath.mj_quatprod([0., 0., 1., 0.], [0., 0., 0., 1.]),
+        [0., 1., 0., 0.])
+    np.testing.assert_allclose(
+        mjmath.mj_quatprod([0., 0., 0., 1.], [0., 1., 0., 0.]),
+        [0., 0., 1., 0.])
+
+  def testQuat2Vel(self):
+    np.testing.assert_allclose(
+        mjmath.mj_quat2vel([0., 1., 0., 0.], 0.1), [math.pi / 0.1, 0., 0.])
+
+  def testQuatNeg(self):
+    np.testing.assert_allclose(
+        mjmath.mj_quatneg([math.sqrt(0.5), math.sqrt(0.5), 0., 0.]),
+        [math.sqrt(0.5), -math.sqrt(0.5), 0., 0.])
+
+  def testQuatDiff(self):
+    np.testing.assert_allclose(
+        mjmath.mj_quatdiff([0., 1., 0., 0.], [0., 0., 1., 0.]),
+        [0., 0., 0., -1.])
+
+  def testEuler2Quat(self):
+    np.testing.assert_allclose(
+        mjmath.euler2quat(0., 0., 0.), [1., 0., 0., 0.])
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/dm_control/suite/utils/parse_amc.py b/dm_control/suite/utils/parse_amc.py
@@ -21,13 +21,11 @@
 
 import collections
 
-from dm_control.mujoco.wrapper import mjbindings
+from dm_control.mujoco import math as mjmath
 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
 
@@ -104,7 +102,8 @@ def convert(file_name, physics, timestep):
     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),
+            mjmath.mj_quat2vel(
+                mjmath.mj_quatdiff(p_t[3:7], p_tp1[3:7]), timestep),
             (p_tp1[7:]-p_t[7:])/ timestep]
     qvel_list.append(np.concatenate(qvel))
 
@@ -187,65 +186,10 @@ def __call__(self, amc_val):
 
     # 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)
+    qpos_quat = mjmath.euler2quat(amc_val[3], amc_val[4], amc_val[5])
+    qpos_quat = mjmath.mj_quatprod(mjmath.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))
diff --git a/setup.py b/setup.py
@@ -165,7 +165,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.286587932',
+    version='0.0.286594366',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
