Add AppendageFixer to regenerate the appendages entries in the trajectories.

sbohez · copybara-github · commit 652690132eda · 2022-03-30T08:26:22.000-07:00
PiperOrigin-RevId: 438307923
Change-Id: I9b8abeb2d4b87f1518c5e498d9507fb742fba0df
diff --git a/dm_control/locomotion/mocap/loader.py b/dm_control/locomotion/mocap/loader.py
@@ -21,6 +21,7 @@
 from dm_control.composer import variation
 from dm_control.locomotion.mocap import mocap_pb2
 from dm_control.locomotion.mocap import trajectory
+from dm_control.utils import transformations as tr
 import numpy as np
 
 from google.protobuf import descriptor
@@ -232,3 +233,15 @@ def __call__(self, proto, random_state=None):
       for prop_pose in t.props:
         # shift prop position
         self._add_z_offset(prop_pose.position)
+
+
+class AppendageFixer:
+
+  def __call__(self, proto, random_state=None):
+    for t in proto.timesteps:
+      for walker_pose in t.walkers:
+        xpos = np.asarray(walker_pose.position)
+        xquat = np.asarray(walker_pose.quaternion)
+        appendages = np.reshape(walker_pose.appendages, (-1, 3))
+        xmat = tr.quat_to_mat(xquat)[:3, :3]
+        walker_pose.appendages[:] = np.ravel((appendages - xpos) @ xmat)
