Expose mj_contactForce() as a method of Physics.data.

yuvaltassa · alimuldal · commit 7de8ff739690 · 2020-05-20T16:39:05.000+01:00
PiperOrigin-RevId: 311731186
Change-Id: Ia8028e12d51d4b14021044ec6048f847f1420259
diff --git a/dm_control/mujoco/wrapper/core.py b/dm_control/mujoco/wrapper/core.py
@@ -52,6 +52,10 @@
 _INVALID_FONT_SCALE = ("`font_scale` must be one of {}, got {{}}."
                        .format(enums.mjtFontScale))
 
+_CONTACT_ID_OUT_OF_RANGE = (
+    "`contact_id` must be between 0 and {max_valid} (inclusive), got: {actual}."
+)
+
 # Global cache used to store finalizers for freeing ctypes pointers.
 # Contains {pointer_address: weakref_object} pairs.
 _FINALIZERS = {}
@@ -714,6 +718,27 @@ def object_velocity(self, object_id, object_type, local_frame=False):
     #  MuJoCo returns velocities in (angular, linear) order, which we flip here.
     return velocity.reshape(2, 3)[::-1]
 
+  def contact_force(self, contact_id):
+    """Returns the wrench of a contact as a 2 x 3 array of (forces, torques).
+
+    Args:
+      contact_id: Integer, the index of the contact within the contact buffer
+        (`self.contact`).
+
+    Returns:
+      2x3 array with stacked (force, torque). Note that the order of dimensions
+        is (normal, tangent, tangent), in the contact's frame.
+
+    Raises:
+      ValueError: If `contact_id` is negative or bigger than ncon-1.
+    """
+    if not 0 <= contact_id < self.ncon:
+      raise ValueError(_CONTACT_ID_OUT_OF_RANGE
+                       .format(max_valid=self.ncon-1, actual=contact_id))
+    wrench = np.empty(6, dtype=np.float64)
+    mjlib.mj_contactForce(self.model.ptr, self.ptr, contact_id, wrench)
+    return wrench.reshape(2, 3)
+
   @property
   def model(self):
     """The parent MjModel for this MjData instance."""
diff --git a/dm_control/mujoco/wrapper/core_test.py b/dm_control/mujoco/wrapper/core_test.py
@@ -465,6 +465,77 @@ def testObjectVelocity(
     np.testing.assert_array_almost_equal(linvel, expected_linvel)
     np.testing.assert_array_almost_equal(angvel, expected_angvel)
 
+  def testContactForce(self):
+    box_on_floor = """
+    <mujoco>
+      <worldbody>
+        <geom name='floor' type='plane' size='1 1 1'/>
+        <body name='box' pos='0 0 .1'>
+          <freejoint/>
+          <geom name='box' type='box' size='.1 .1 .1'/>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+    model = core.MjModel.from_xml_string(box_on_floor)
+    data = core.MjData(model)
+    # Settle for 500 timesteps (1 second):
+    for _ in range(500):
+      mjlib.mj_step(model.ptr, data.ptr)
+    normal_force = 0.
+    for contact_id in range(data.ncon):
+      force = data.contact_force(contact_id)
+      normal_force += force[0, 0]
+    box_id = 1
+    box_weight = -model.opt.gravity[2]*model.body_mass[box_id]
+    self.assertAlmostEqual(normal_force, box_weight)
+    # Test raising of out-of-range errors:
+    bad_ids = [-1, data.ncon]
+    for bad_id in bad_ids:
+      with self.assertRaisesWithLiteralMatch(
+          ValueError,
+          core._CONTACT_ID_OUT_OF_RANGE.format(
+              max_valid=data.ncon - 1, actual=bad_id)):
+        data.contact_force(bad_id)
+
+  @parameterized.parameters(
+      dict(
+          condim=3,  # Only sliding friction.
+          expected_torques=[False, False, False],  # No torques.
+      ),
+      dict(
+          condim=4,  # Sliding and torsional friction.
+          expected_torques=[True, False, False],  # Only torsional torque.
+      ),
+      dict(
+          condim=6,  # Sliding, torsional and rolling.
+          expected_torques=[True, True, True],  # All torques are nonzero.
+      ),
+  )
+  def testContactTorque(self, condim, expected_torques):
+    ball_on_floor = """
+    <mujoco>
+      <worldbody>
+        <geom name='floor' type='plane' size='1 1 1'/>
+        <body name='ball' pos='0 0 .1'>
+          <freejoint/>
+          <geom name='ball' size='.1' friction='1 .1 .1'/>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+    model = core.MjModel.from_xml_string(ball_on_floor)
+    data = core.MjData(model)
+    model.geom_condim[:] = condim
+    data.qvel[3:] = np.array((1., 1., 1.))
+    # Settle for 10 timesteps (20 milliseconds):
+    for _ in range(10):
+      mjlib.mj_step(model.ptr, data.ptr)
+    contact_id = 0  # This model has only one contact.
+    _, torque = data.contact_force(contact_id)
+    nonzero_torques = torque != 0
+    np.testing.assert_array_equal(nonzero_torques, np.array((expected_torques)))
+
   def testFreeMjrContext(self):
     for _ in range(5):
       renderer = _render.Renderer(640, 480)
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.311497959',
+    version='0.0.311731186',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
