Move implementation of transform_{vec,xmat}_to_egocentric_frame from base.Walker to base.Entity.

saran-t · copybara-github · commit b5a1c4e17234 · 2020-02-24T08:49:27.000-08:00
New methods in `base.Entity` are renamed to `global_{vector,xmat}_to_local_frame`.

PiperOrigin-RevId: 296900770
Change-Id: I99538d7c4e72661d179ddcd4a84bf8464357c99b
diff --git a/dm_control/composer/entity.py b/dm_control/composer/entity.py
@@ -335,6 +335,84 @@ def parent(self):
   def attachment_site(self):
     return self.mjcf_model
 
+  @property
+  def root_body(self):
+    if self.parent:
+      return mjcf.get_attachment_frame(self.mjcf_model)
+    else:
+      return self.mjcf_model.worldbody
+
+  def global_vector_to_local_frame(self, physics, vec_in_world_frame):
+    """Linearly transforms a world-frame vector into entity's local frame.
+
+    Note that this function does not perform an affine transformation of the
+    vector. In other words, the input vector is assumed to be specified with
+    respect to the same origin as this entity's local frame. This function
+    can also be applied to matrices whose innermost dimensions are either 2 or
+    3. In this case, a matrix with the same leading dimensions is returned
+    where the innermost vectors are replaced by their values computed in the
+    local frame.
+
+    Args:
+      physics: An `mjcf.Physics` instance.
+      vec_in_world_frame: A NumPy array with last dimension of shape (2,) or
+      (3,) that represents a vector quantity in the world frame.
+
+    Returns:
+      The same quantity as `vec_in_world_frame` but reexpressed in this
+      entity's local frame. The returned np.array has the same shape as
+      np.asarray(vec_in_world_frame).
+
+    Raises:
+      ValueError: if `vec_in_world_frame` does not have shape ending with (2,)
+        or (3,).
+    """
+    vec_in_world_frame = np.asarray(vec_in_world_frame)
+
+    xmat = np.reshape(physics.bind(self.root_body).xmat, (3, 3))
+    # The ordering of the np.dot is such that the transformation holds for any
+    # matrix whose final dimensions are (2,) or (3,).
+    if vec_in_world_frame.shape[-1] == 2:
+      return np.dot(vec_in_world_frame, xmat[:2, :2])
+    elif vec_in_world_frame.shape[-1] == 3:
+      return np.dot(vec_in_world_frame, xmat)
+    else:
+      raise ValueError('`vec_in_world_frame` should have shape with final '
+                       'dimension 2 or 3: got {}'.format(
+                           vec_in_world_frame.shape))
+
+  def global_xmat_to_local_frame(self, physics, xmat):
+    """Transforms another entity's `xmat` into this entity's local frame.
+
+    This function takes another entity's (E) xmat, which is an SO(3) matrix
+    from E's frame to the world frame, and turns it to a matrix that transforms
+    from E's frame into this entity's local frame.
+
+    Args:
+      physics: An `mjcf.Physics` instance.
+      xmat: A NumPy array of shape (3, 3) or (9,) that represents another
+        entity's xmat.
+
+    Returns:
+      The `xmat` reexpressed in this entity's local frame. The returned
+      np.array has the same shape as np.asarray(xmat).
+
+    Raises:
+      ValueError: if `xmat` does not have shape (3, 3) or (9,).
+    """
+    xmat = np.asarray(xmat)
+
+    input_shape = xmat.shape
+    if xmat.shape == (9,):
+      xmat = np.reshape(xmat, (3, 3))
+
+    self_xmat = np.reshape(physics.bind(self.root_body).xmat, (3, 3))
+    if xmat.shape == (3, 3):
+      return np.reshape(np.dot(self_xmat.T, xmat), input_shape)
+    else:
+      raise ValueError('`xmat` should have shape (3, 3) or (9,): got {}'.format(
+          xmat.shape))
+
   def get_pose(self, physics):
     """Get the position and orientation of this entity relative to its parent.
 
diff --git a/dm_control/composer/entity_test.py b/dm_control/composer/entity_test.py
@@ -292,6 +292,51 @@ def testIterEntitiesExcludeSelf(self):
     self.assertEqual(
         list(entities[0].iter_entities(exclude_self=True)), entities[1:])
 
+  def testGlobalVectorToLocalFrame(self):
+    parent = TestEntity()
+    parent.mjcf_model.worldbody.add(
+        'site', xyaxes=[0, 1, 0, -1, 0, 0]).attach(self.entity.mjcf_model)
+    physics = mjcf.Physics.from_mjcf_model(parent.mjcf_model)
+
+    # 3D vectors
+    np.testing.assert_allclose(
+        self.entity.global_vector_to_local_frame(physics, [0, 1, 0]),
+        [1, 0, 0], atol=1e-10)
+    np.testing.assert_allclose(
+        self.entity.global_vector_to_local_frame(physics, [-1, 0, 0]),
+        [0, 1, 0], atol=1e-10)
+    np.testing.assert_allclose(
+        self.entity.global_vector_to_local_frame(physics, [0, 0, 1]),
+        [0, 0, 1], atol=1e-10)
+
+    # 2D vectors; z-component is ignored
+    np.testing.assert_allclose(
+        self.entity.global_vector_to_local_frame(physics, [0, 1]),
+        [1, 0], atol=1e-10)
+    np.testing.assert_allclose(
+        self.entity.global_vector_to_local_frame(physics, [-1, 0]),
+        [0, 1], atol=1e-10)
+
+  def testGlobalMatrixToLocalFrame(self):
+    parent = TestEntity()
+    parent.mjcf_model.worldbody.add(
+        'site', xyaxes=[0, 1, 0, -1, 0, 0]).attach(self.entity.mjcf_model)
+    physics = mjcf.Physics.from_mjcf_model(parent.mjcf_model)
+
+    rotation_atob = np.array([[0, 1, 0], [0, 0, -1], [-1, 0, 0]])
+    ego_rotation_atob = np.array([[0, 0, -1], [0, -1, 0], [-1, 0, 0]])
+
+    np.testing.assert_allclose(
+        self.entity.global_xmat_to_local_frame(physics, rotation_atob),
+        ego_rotation_atob, atol=1e-10)
+
+    flat_rotation_atob = np.reshape(rotation_atob, -1)
+    flat_rotation_ego_atob = np.reshape(ego_rotation_atob, -1)
+    np.testing.assert_allclose(
+        self.entity.global_xmat_to_local_frame(
+            physics, flat_rotation_atob),
+        flat_rotation_ego_atob, atol=1e-10)
+
   @parameterized.parameters(*_param_product(
       position=[None, [1., 0., -1.]],
       quaternion=[None, _FORTYFIVE_DEGREES_ABOUT_X, _NINETY_DEGREES_ABOUT_Z],
diff --git a/dm_control/locomotion/walkers/base.py b/dm_control/locomotion/walkers/base.py
@@ -105,19 +105,8 @@ def transform_vec_to_egocentric_frame(self, physics, vec_in_world_frame):
       ValueError: if `vec_in_world_frame` does not have shape ending with (2,)
         or (3,).
     """
-    vec_in_world_frame = np.asarray(vec_in_world_frame)
-
-    xmat = np.reshape(physics.bind(self.root_body).xmat, (3, 3))
-    # The ordering of the np.dot is such that the transformation holds for any
-    # matrix whose final dimensions are (2,) or (3,).
-    if vec_in_world_frame.shape[-1] == 2:
-      return np.dot(vec_in_world_frame, xmat[:2, :2])
-    elif vec_in_world_frame.shape[-1] == 3:
-      return np.dot(vec_in_world_frame, xmat)
-    else:
-      raise ValueError('`vec_in_world_frame` should have shape with final '
-                       'dimension 2 or 3: got {}'.format(
-                           vec_in_world_frame.shape))
+    return super(Walker, self).global_vector_to_local_frame(
+        physics, vec_in_world_frame)
 
   def transform_xmat_to_egocentric_frame(self, physics, xmat):
     """Transforms another entity's `xmat` into this walker's egocentric frame.
@@ -138,18 +127,7 @@ def transform_xmat_to_egocentric_frame(self, physics, xmat):
     Raises:
       ValueError: if `xmat` does not have shape (3, 3) or (9,).
     """
-    xmat = np.asarray(xmat)
-
-    input_shape = xmat.shape
-    if xmat.shape == (9,):
-      xmat = np.reshape(xmat, (3, 3))
-
-    self_xmat = np.reshape(physics.bind(self.root_body).xmat, (3, 3))
-    if xmat.shape == (3, 3):
-      return np.reshape(np.dot(self_xmat.T, xmat), input_shape)
-    else:
-      raise ValueError('`xmat` should have shape (3, 3) or (9,): got {}'.format(
-          xmat.shape))
+    return super(Walker, self).global_xmat_to_local_frame(physics, xmat)
 
   @abc.abstractproperty
   def root_body(self):
diff --git a/setup.py b/setup.py
@@ -166,7 +166,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.296467089',
+    version='0.0.296900770',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
