Add QuatRotate variation and use dm_control.utils.transformations for calcs.

PiperOrigin-RevId: 719432052
Change-Id: I17cfbe510532b6b97212bd1c5ae87f5aa204ee78

Author: sbohez

dm_control/composer/variation/rotations.py

@@ -18,6 +18,7 @@
 
 from dm_control.composer.variation import base
 from dm_control.composer.variation import variation_values
+from dm_control.utils import transformations
 import numpy as np
 
 IDENTITY_QUATERNION = np.array([1., 0., 0., 0.])
@@ -48,8 +49,7 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
         self._axis, initial_value, current_value, random_state)
     angle = variation_values.evaluate(
         self._angle, initial_value, current_value, random_state)
-    sine, cosine = np.sin(angle / 2), np.cos(angle / 2)
-    return np.array([cosine, axis[0] * sine, axis[1] * sine, axis[2] * sine])
+    return transformations.axisangle_to_quat(np.asarray(axis) * angle)
 
 
 class QuaternionPreMultiply(base.Variation):
@@ -70,8 +70,32 @@ def __call__(self, initial_value=None, current_value=None, random_state=None):
     q1 = variation_values.evaluate(self._quat, initial_value, current_value,
                                    random_state)
     q2 = current_value if self._cumulative else initial_value
-    return np.array([
-        q1[0]*q2[0] - q1[1]*q2[1] - q1[2]*q2[2] - q1[3]*q2[3],
-        q1[0]*q2[1] + q1[1]*q2[0] + q1[2]*q2[3] - q1[3]*q2[2],
-        q1[0]*q2[2] - q1[1]*q2[3] + q1[2]*q2[0] + q1[3]*q2[1],
-        q1[0]*q2[3] + q1[1]*q2[2] - q1[2]*q2[1] + q1[3]*q2[0]])
+    return transformations.quat_mul(np.asarray(q1), np.asarray(q2))
+
+
+class QuaternionRotate(base.Variation):
+  """Variation that rotates a given vector by the given quaternion.
+
+  The vector can either be an existing value passed at evaluation, or specified
+  as a separate variation at construction. In the former case, cumulative mode
+  determines whether to use the current or initial value of the vector. The#
+  quaternion is always specified by a variation at construction.
+  """
+
+  def __init__(self, quat, vec=None, cumulative=False):
+    self._quat = quat
+    self._vec = vec
+    self._cumulative = cumulative
+
+  def __call__(self, initial_value=None, current_value=None, random_state=None):
+    random_state = random_state or np.random
+    quat = variation_values.evaluate(
+        self._quat, initial_value, current_value, random_state
+    )
+    if self._vec is None:
+      vec = current_value if self._cumulative else initial_value
+    else:
+      vec = variation_values.evaluate(
+          self._vec, initial_value, current_value, random_state
+      )
+    return transformations.quat_rotate(np.asarray(quat), np.asarray(vec))

dm_control/utils/transformations.py

@@ -625,6 +625,26 @@ def mat_to_quat(mat):
   return q
 
 
+def axisangle_to_quat(axisangle, tol=0.0):
+  """Returns the quaternion corresponding to the provided axis-angle.
+
+  Args:
+    axisangle: A 3x1 numpy array describing the axis of rotation, with angle
+      encoded by its length.
+    tol: Tolerance for the angle magnitude below which the identity quaternion
+      is returned.
+
+  Returns:
+    A quaternion [w, i, j, k].
+  """
+  axisangle = np.asarray(axisangle)
+  angle = np.linalg.norm(axisangle, axis=-1, keepdims=True)
+  axis = np.where(angle <= tol, [1.0, 0.0, 0.0], axisangle / angle)
+  angle = np.where(angle <= tol, [0.0], angle)
+  sine, cosine = np.sin(angle / 2), np.cos(angle / 2)
+  return np.concatenate([cosine, axis * sine], axis=-1)
+
+
 # ################
 # # 2D Functions #
 # ################

dm_control/utils/transformations_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ============================================================================
 
-"""Tests for dm_control.locomotion.tasks.transformations."""
-
 import itertools
 
 from absl.testing import absltest
@@ -261,6 +259,35 @@ def _random_quaternion(self):
         (np.cos(t2) * r2, np.sin(t1) * r1, np.cos(t1) * r1, np.sin(t2) * r2),
         dtype=np.float64)
 
+  def test_axisangle_to_quat(self):
+    axisangle = np.array([0.1, 0.2, 0.3])
+    quat = transformations.axisangle_to_quat(axisangle)
+    np.testing.assert_allclose(
+        quat, [0.982551, 0.0497088, 0.0994177, 0.1491265], atol=1e-6
+    )
+
+  def test_axisangle_to_quat_zero(self):
+    axisangle = np.array([0, 0, 0])
+    quat = transformations.axisangle_to_quat(axisangle)
+    np.testing.assert_allclose(quat, [1, 0, 0, 0])
+
+  def test_axisangle_to_quat_zero_tol(self):
+    axisangle = np.array([0, 0, 1e-2])
+    quat = transformations.axisangle_to_quat(axisangle, tol=1e-1)
+    np.testing.assert_allclose(quat, [1, 0, 0, 0])
+
+  def test_axisangle_to_quat_batched(self):
+    axisangle = np.stack([np.array([0.1, 0.2, 0.3]), np.array([0.4, 0.5, 0.6])])
+    quat = transformations.axisangle_to_quat(axisangle)
+    np.testing.assert_allclose(
+        quat,
+        [
+            [0.982551, 0.0497088, 0.0994177, 0.1491265],
+            [0.9052841, 0.1936448, 0.242056, 0.2904672],
+        ],
+        atol=1e-6,
+    )
+
 
 if __name__ == '__main__':
   absltest.main()