Locomotion reference pose task: Add more general quaternion difference reward computation.

DeepMind · alimuldal · commit db36290c588f · 2021-04-26T10:47:03.000+01:00
PiperOrigin-RevId: 369839626
Change-Id: I19a67a299e06143ba8fca86b6a00b1aff51f53fb
diff --git a/dm_control/locomotion/tasks/reference_pose/rewards.py b/dm_control/locomotion/tasks/reference_pose/rewards.py
@@ -36,6 +36,8 @@ def bounded_quat_dist(source, target):
   Returns:
     quaternion distance.
   """
+  source /= np.linalg.norm(source)
+  target /= np.linalg.norm(target)
   default_dist = tr.quat_dist(source, target)
   anti_dist = tr.quat_dist(source, -np.asarray(target))
   min_dist = np.minimum(default_dist, anti_dist)
@@ -55,14 +57,16 @@ def compute_squared_differences(walker_features, reference_features,
       if 'quaternion' not in k:
         squared_differences[k] = np.sum(
             (walker_features[k] - reference_features[k])**2)
-  quat_dists = np.array([
-      bounded_quat_dist(w, r)
-      for w, r in zip(walker_features['body_quaternions'],
-                      reference_features['body_quaternions'])
-  ])
-  squared_differences['body_quaternions'] = np.sum(quat_dists**2)
-  squared_differences['quaternion'] = bounded_quat_dist(
-      walker_features['quaternion'], reference_features['quaternion'])**2
+      elif 'quaternions' in k:
+        quat_dists = np.array([
+            bounded_quat_dist(w, r)
+            for w, r in zip(walker_features[k], reference_features[k])
+        ])
+        squared_differences[k] = np.sum(quat_dists**2)
+      else:
+        squared_differences[k] = bounded_quat_dist(walker_features[k],
+                                                   reference_features[k])**2
+
   return squared_differences
 
 
diff --git a/dm_control/locomotion/tasks/reference_pose/rewards_test.py b/dm_control/locomotion/tasks/reference_pose/rewards_test.py
@@ -0,0 +1,84 @@
+# Copyright 2021 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 dm_control.locomotion.tasks.reference_pose.rewards."""
+
+from absl.testing import absltest
+from dm_control.locomotion.tasks.reference_pose import rewards
+import numpy as np
+
+WALKER_FEATURES = {
+    'scalar': 0.,
+    'vector': np.ones(3),
+    'match': 0.1,
+}
+
+REFERENCE_FEATURES = {
+    'scalar': 1.5,
+    'vector': np.full(3, 2),
+    'match': 0.1,
+}
+
+QUATERNION_FEATURES = {
+    'unmatched_quaternion': (1., 0., 0., 0.),
+    'matched_quaternions': [(1., 0., 1., 0.), (0.707, 0.707, 0., 0.)],
+}
+
+REFERENCE_QUATERNION_FEATURES = {
+    'unmatched_quaternion': (0., 0., 0., 1.),
+    'matched_quaternions': [(1., 0., 1., 0.), (0.707, 0.707, 0., 0.)],
+}
+
+
+EXPECTED_DIFFERENCES = {
+    'scalar': 2.25,
+    'vector': 3.,
+    'match': 0.,
+    'unmatched_quaternion': np.sum(rewards.bounded_quat_dist(
+        QUATERNION_FEATURES['unmatched_quaternion'],
+        REFERENCE_QUATERNION_FEATURES['unmatched_quaternion']))**2,
+    'matched_quaternions': 0.,
+}
+
+EXCLUDE_KEYS = ('scalar', 'match')
+
+
+class RewardsTest(absltest.TestCase):
+
+  def test_compute_squared_differences(self):
+    """Basic usage."""
+    differences = rewards.compute_squared_differences(
+        WALKER_FEATURES, REFERENCE_FEATURES)
+    for key, difference in differences.items():
+      self.assertEqual(difference, EXPECTED_DIFFERENCES[key])
+
+  def test_compute_squared_differences_exclude_keys(self):
+    """Test excluding some keys from squared difference computation."""
+    differences = rewards.compute_squared_differences(
+        WALKER_FEATURES, REFERENCE_FEATURES, exclude_keys=EXCLUDE_KEYS)
+    for key in EXCLUDE_KEYS:
+      self.assertNotIn(key, differences)
+
+  def test_compute_squared_differences_quaternion(self):
+    """Test that quaternions use a different distance computation."""
+
+    differences = rewards.compute_squared_differences(
+        QUATERNION_FEATURES, REFERENCE_QUATERNION_FEATURES)
+
+    for key, difference in differences.items():
+      self.assertAlmostEqual(difference, EXPECTED_DIFFERENCES[key])
+
+
+if __name__ == '__main__':
+  absltest.main()
