# Copyright 2017 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.
# ============================================================================

import copy
import pickle

from absl.testing import absltest
from absl.testing import parameterized
from dm_control.mujoco import engine
from dm_control.mujoco import wrapper
from dm_control.mujoco.testing import assets
from dm_control.mujoco.wrapper.mjbindings import enums
from dm_control.rl import control
import mock
import mujoco
import numpy as np

MODEL_PATH = assets.get_path('cartpole.xml')
MODEL_WITH_ASSETS = assets.get_contents('model_with_assets.xml')
ASSETS = {
    'texture.png': assets.get_contents('deepmind.png'),
    'mesh.stl': assets.get_contents('cube.stl'),
    'included.xml': assets.get_contents('sphere.xml')
}


class MujocoEngineTest(parameterized.TestCase):

  def setUp(self):
    super().setUp()
    self._physics = engine.Physics.from_xml_path(MODEL_PATH)

  def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
    for name in attr_to_compare:
      actual_value = getattr(actual_obj, name)
      expected_value = getattr(expected_obj, name)
      try:
        if isinstance(expected_value, np.ndarray):
          np.testing.assert_array_equal(actual_value, expected_value)
        else:
          self.assertEqual(actual_value, expected_value)
      except AssertionError as e:
        raise AssertionError(
            f"Attribute '{name}' differs from expected value.") from e

  @parameterized.parameters(0, 'cart', u'cart')
  def testCameraIndexing(self, camera_id):
    height, width = 480, 640
    _ = engine.Camera(
        self._physics, height, width, camera_id=camera_id)

  def testDepthRender(self):
    plane_and_box = """
    <mujoco>
      <worldbody>
        <geom type="plane" pos="0 0 0" size="2 2 .1"/>
        <geom type="box" size=".1 .1 .1" pos="0 0 .1"/>
        <camera name="top" pos="0 0 3"/>
      </worldbody>
    </mujoco>
    """
    physics = engine.Physics.from_xml_string(plane_and_box)
    pixels = physics.render(height=200, width=200, camera_id='top', depth=True)
    # Nearest pixels should be 2.8m away
    np.testing.assert_approx_equal(pixels.min(), 2.8, 3)
    # Furthest pixels should be 3m away (depth is orthographic)
    np.testing.assert_approx_equal(pixels.max(), 3.0, 3)

  @parameterized.parameters([True, False])
  def testSegmentationRender(self, enable_geom_frame_rendering):
    box_four_corners = """
    <mujoco>
      <visual>
        <scale framelength="2"/>
      </visual>
      <worldbody>
        <geom name="box0" type="box" size=".2 .2 .2" pos="-1 1 .1"/>
        <geom name="box1" type="box" size=".2 .2 .2" pos="1 1 .1"/>
        <site name="box2" type="box" size=".2 .2 .2" pos="1 -1 .1"/>
        <site name="box3" type="box" size=".2 .2 .2" pos="-1 -1 .1"/>
        <camera name="top" pos="0 0 3"/>
      </worldbody>
    </mujoco>
    """
    physics = engine.Physics.from_xml_string(box_four_corners)
    obj_type_geom = enums.mjtObj.mjOBJ_GEOM  # Geom object type
    obj_type_site = enums.mjtObj.mjOBJ_SITE  # Site object type
    obj_type_decor = enums.mjtObj.mjOBJ_UNKNOWN  # Decor object type
    scene_options = wrapper.MjvOption()
    if enable_geom_frame_rendering:
      scene_options.frame = mujoco.mjtFrame.mjFRAME_GEOM
    pixels = physics.render(height=200, width=200, camera_id='top',
                            segmentation=True, scene_option=scene_options)

    # The pixel indices below were chosen so that toggling the frame decors do
    # not affect the segmentation results.
    with self.subTest('Center pixels should have background label'):
      np.testing.assert_equal(pixels[95:105, 95:105, 0], -1)
      np.testing.assert_equal(pixels[95:105, 95:105, 1], -1)
    with self.subTest('Geoms have correct object type'):
      np.testing.assert_equal(pixels[15:25, 0:10, 1], obj_type_geom)
      np.testing.assert_equal(pixels[15:25, 190:200, 1], obj_type_geom)
    with self.subTest('Sites have correct object type'):
      np.testing.assert_equal(pixels[190:200, 190:200, 1], obj_type_site)
      np.testing.assert_equal(pixels[190:200, 0:10, 1], obj_type_site)
    with self.subTest('Geoms have correct object IDs'):
      np.testing.assert_equal(pixels[15:25, 0:10, 0],
                              physics.model.name2id('box0', obj_type_geom))
      np.testing.assert_equal(pixels[15:25, 190:200, 0],
                              physics.model.name2id('box1', obj_type_geom))
    with self.subTest('Sites have correct object IDs'):
      np.testing.assert_equal(pixels[190:200, 190:200, 0],
                              physics.model.name2id('box2', obj_type_site))
      np.testing.assert_equal(pixels[190:200, 0:10, 0],
                              physics.model.name2id('box3', obj_type_site))
    with self.subTest('Decor elements present if and only if geom frames are '
                      'enabled'):
      contains_decor = np.any(pixels[:, :, 1] == obj_type_decor)
      self.assertEqual(contains_decor, enable_geom_frame_rendering)

  def testSceneCallback(self):
    empty_world = """
    <mujoco>
      <worldbody>
        <camera name="cam" pos="0 0 3"/>
      </worldbody>
    </mujoco>
    """

    def callback(_, scn: mujoco.MjvScene):
      # Add a red box to the scene
      scn.ngeom += 1
      mujoco.mjv_initGeom(
          scn.geoms[scn.ngeom - 1],
          mujoco.mjtGeom.mjGEOM_BOX.value,
          size=np.array([0.2, 0.2, 0.2]),
          pos=np.zeros(3),
          mat=np.eye(3).flatten(),
          rgba=np.array([1, 0, 0, 1], dtype=np.float32))

    physics = engine.Physics.from_xml_string(empty_world)

    # Without the callback, render should return a black image.
    empty_image = physics.render(
        height=8, width=8, camera_id='cam', scene_callback=None)
    np.testing.assert_array_equal(
        np.zeros((8, 8, 3), dtype=np.uint8), empty_image)

    # With the callback, there should be a red box.
    pixels = physics.render(
        height=8, width=8, camera_id='cam', scene_callback=callback)
    # Are there any pixels where red component is bigger than green and blue?
    any_red_pixels = np.any(pixels[:, :, 0] > np.max(pixels[:, :, 1:3], axis=2))
    self.assertTrue(any_red_pixels, 'Expecting some red pixels.')

  def testTextOverlay(self):
    height, width = 480, 640
    overlay = engine.TextOverlay(title='Title', body='Body', style='big',
                                 position='bottom right')

    no_overlay = self._physics.render(height, width, camera_id=0)
    with_overlay = self._physics.render(height, width, camera_id=0,
                                        overlays=[overlay])
    self.assertFalse(np.all(no_overlay == with_overlay),
                     msg='Images are identical with and without text overlay.')

  def testSceneOption(self):
    height, width = 480, 640
    scene_option = wrapper.MjvOption()

    # Render geoms as semi-transparent.
    scene_option.flags[enums.mjtVisFlag.mjVIS_TRANSPARENT] = 1

    no_scene_option = self._physics.render(height, width, camera_id=0)
    with_scene_option = self._physics.render(height, width, camera_id=0,
                                             scene_option=scene_option)
    self.assertFalse(np.all(no_scene_option == with_scene_option),
                     msg='Images are identical with and without scene option.')

  def testRenderFlags(self):
    height, width = 480, 640
    cam = engine.Camera(self._physics, height, width, camera_id=0)
    cam.scene.flags[enums.mjtRndFlag.mjRND_WIREFRAME] = 1  # Enable wireframe
    enabled = cam.render().copy()
    cam.scene.flags[enums.mjtRndFlag.mjRND_WIREFRAME] = 0  # Disable wireframe
    disabled = cam.render().copy()
    self.assertFalse(
        np.all(disabled == enabled),
        msg='Images are identical regardless of whether wireframe is enabled.')

  def testFreeIdempotent(self):
    self._physics.free()
    self._physics.free()

  @parameterized.parameters(((0.5, 0.5), (1, 3)),  # pole
                            ((0.5, 0.1), (0, 0)),  # ground
                            ((0.9, 0.9), (None, None)),  # sky
                           )
  def testCameraSelection(self, coordinates, expected_selection):
    height, width = 480, 640
    camera = engine.Camera(self._physics, height, width, camera_id=0)

    # Test for b/63380170: Enabling visualization of body frames adds
    # "non-model" geoms to the scene. This means that the indices of geoms
    # within `camera._scene.geoms` don't match the rows of `model.geom_bodyid`.
    camera.option.frame = enums.mjtFrame.mjFRAME_BODY

    selected = camera.select(coordinates)
    self.assertEqual(expected_selection, selected[:2])

  @parameterized.parameters(
      dict(camera_id='cam0', height=200, width=300),
      dict(camera_id=1, height=300, width=200),
      dict(camera_id=-1, height=400, width=400),
  )
  def testCameraMatrix(self, camera_id, height, width):
    """Tests the camera_matrix() method.

       Creates a model with two cameras and two small geoms. We render the scene
       with one of the cameras and check that the geom locations, projected into
       pixel space, are correct, using segmenation rendering.
       xyz2pixels() shows how the transformation is used. For a description
       of the camera matrix see https://en.wikipedia.org/wiki/Camera_matrix.

    Args:
      camera_id: One of the two cameras. Can be either integer or String.
      height: The height of the image (pixels).
      width: The width of the image (pixels).
    """

    def xyz2pixels(x, y, z, camera_matrix):
      """Transforms from world coordinates to pixel coordinates."""
      xs, ys, s = camera_matrix.dot(np.array([x, y, z, 1.0]))
      return xs/s, ys/s

    two_geoms_and_two_cameras = """
    <mujoco>
      <visual>
        <global fovy="55"/>
      </visual>
      <worldbody>
        <light name="top" pos="0 0 1"/>
        <geom name="red" pos=".2 0 0" size=".005" rgba="1 0 0 1"/>
        <geom name="green" pos=".2 .2 .1" size=".005" rgba="0 1 0 1"/>
        <camera name="cam0" pos="1 .5 1" zaxis="1 .5 1" fovy="20"/>
        <camera name="cam1" pos=".1 .1 1" xyaxes="1 1 0 -1 0 0"/>
      </worldbody>
    </mujoco>
    """
    physics = engine.Physics.from_xml_string(two_geoms_and_two_cameras)
    camera = engine.Camera(physics, width=width, height=height,
                           camera_id=camera_id)
    camera_matrix = camera.matrix  # Get camera matrix.
    pixels = camera.render(segmentation=True)  # Render a segmentation frame.
    for geom_id in [0, 1]:
      # Compute the location of the geom in pixel space using the camera matrix.
      x, y = xyz2pixels(*physics.data.geom_xpos[geom_id], camera_matrix)
      row = int(round(y))
      column = int(round(x))
      # Compare segmentation values of nearest pixel to corresponding geom.
      [obj_id, obj_type] = pixels[row, column, :]
      self.assertEqual(obj_type, enums.mjtObj.mjOBJ_GEOM)
      self.assertEqual(obj_id, geom_id)

  def testMovableCameraSetGetPose(self):
    height, width = 240, 320

    camera = engine.MovableCamera(self._physics, height, width)
    image = camera.render().copy()

    pose = camera.get_pose()

    lookat_offset = np.array([0.01, 0.02, -0.03])

    # Would normally pass the new values directly to camera.set_pose instead of
    # using the namedtuple _replace method, but this makes the asserts at the
    # end of the test a little cleaner.
    new_pose = pose._replace(distance=pose.distance * 1.5,
                             lookat=pose.lookat + lookat_offset,
                             azimuth=pose.azimuth + -15,
                             elevation=pose.elevation - 10)

    camera.set_pose(*new_pose)

    self.assertEqual(new_pose.distance, camera.get_pose().distance)
    self.assertEqual(new_pose.azimuth, camera.get_pose().azimuth)
    self.assertEqual(new_pose.elevation, camera.get_pose().elevation)
    np.testing.assert_allclose(new_pose.lookat, camera.get_pose().lookat)

    self.assertFalse(np.all(image == camera.render()))

  def testRenderExceptions(self):
    max_width = self._physics.model.vis.global_.offwidth
    max_height = self._physics.model.vis.global_.offheight
    max_camid = self._physics.model.ncam - 1
    with self.assertRaisesRegex(ValueError, 'width'):
      self._physics.render(max_height, max_width + 1, camera_id=max_camid)
    with self.assertRaisesRegex(ValueError, 'height'):
      self._physics.render(max_height + 1, max_width, camera_id=max_camid)
    with self.assertRaisesRegex(ValueError, 'camera_id'):
      self._physics.render(max_height, max_width, camera_id=max_camid + 1)
    with self.assertRaisesRegex(ValueError, 'camera_id'):
      self._physics.render(max_height, max_width, camera_id=-2)

  def testPhysicsRenderMethod(self):
    height, width = 240, 320
    image = self._physics.render(height=height, width=width)
    self.assertEqual(image.shape, (height, width, 3))
    depth = self._physics.render(height=height, width=width, depth=True)
    self.assertEqual(depth.shape, (height, width))
    segmentation = self._physics.render(height=height, width=width,
                                        segmentation=True)
    self.assertEqual(segmentation.shape, (height, width, 2))

  def testExceptionIfBothDepthAndSegmentation(self):
    with self.assertRaisesWithLiteralMatch(
        ValueError, engine._BOTH_SEGMENTATION_AND_DEPTH_ENABLED):
      self._physics.render(depth=True, segmentation=True)

  def testRenderFlagOverridesAreNotPersistent(self):
    camera = engine.Camera(self._physics)
    first_rgb = camera.render().copy()
    camera.render(segmentation=True)
    second_rgb = camera.render().copy()
    np.testing.assert_array_equal(first_rgb, second_rgb)

  def testCustomRenderFlags(self):
    default = self._physics.render()
    wireframe_string_key = self._physics.render(
        render_flag_overrides=dict(wireframe=True))
    self.assertFalse((default == wireframe_string_key).all())
    wireframe_enum_key = self._physics.render(
        render_flag_overrides={enums.mjtRndFlag.mjRND_WIREFRAME: True})
    np.testing.assert_array_equal(wireframe_string_key, wireframe_enum_key)

  @parameterized.parameters(dict(depth=True), dict(segmentation=True))
  def testExceptionIfRenderFlagOverridesAndDepthOrSegmentation(self, **kwargs):
    with self.assertRaisesWithLiteralMatch(
        ValueError,
        engine._RENDER_FLAG_OVERRIDES_NOT_SUPPORTED_FOR_DEPTH_OR_SEGMENTATION):
      self._physics.render(render_flag_overrides=dict(wireframe=True), **kwargs)

  def testExceptionIfOverlaysAndDepthOrSegmentation(self):
    overlay = engine.TextOverlay()
    with self.assertRaisesWithLiteralMatch(
        ValueError, engine._OVERLAYS_NOT_SUPPORTED_FOR_DEPTH_OR_SEGMENTATION):
      self._physics.render(depth=True, overlays=[overlay])
    with self.assertRaisesWithLiteralMatch(
        ValueError, engine._OVERLAYS_NOT_SUPPORTED_FOR_DEPTH_OR_SEGMENTATION):
      self._physics.render(segmentation=True, overlays=[overlay])

  def testNamedViews(self):
    self.assertEqual((1,), self._physics.control().shape)
    self.assertEqual((2,), self._physics.position().shape)
    self.assertEqual((2,), self._physics.velocity().shape)
    self.assertEqual((0,), self._physics.activation().shape)
    self.assertEqual((4,), self._physics.state().shape)
    self.assertEqual(0., self._physics.time())
    self.assertEqual(0.01, self._physics.timestep())

  def testSetGetPhysicsStateLegacy(self):
    physics_state = self._physics.get_state()

    # qpos, qvel, act
    self.assertLen(self._physics._physics_state_items(), 3)

    self._physics.set_state(physics_state)

    new_physics_state = np.random.random_sample(physics_state.shape)
    self._physics.set_state(new_physics_state)

    np.testing.assert_allclose(new_physics_state,
                               self._physics.get_state())

  def testSetGetPhysicsState(self):
    actual_physics_state = self._physics.get_state(
        sig=mujoco.mjtState.mjSTATE_FULLPHYSICS.value
    )
    expected_physics_state = np.zeros(mujoco.mj_stateSize(
        self._physics.model.ptr, mujoco.mjtState.mjSTATE_FULLPHYSICS.value
    ))
    mujoco.mj_getState(self._physics.model.ptr, self._physics.data.ptr,
                       expected_physics_state,
                       mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
    np.testing.assert_array_equal(expected_physics_state, actual_physics_state)

    new_physics_state = np.random.random_sample(expected_physics_state.shape)
    self._physics.set_state(
        new_physics_state, sig=mujoco.mjtState.mjSTATE_FULLPHYSICS.value
    )
    np.testing.assert_array_equal(
        new_physics_state,
        self._physics.get_state(sig=mujoco.mjtState.mjSTATE_FULLPHYSICS.value),
    )

  def testSetGetPhysicsStateWithPlugin(self):
    # Model copied from mujoco/test/plugin/elasticity/elasticity_test.cc
    model_with_cable_plugin = """
    <mujoco>
      <option gravity="0 0 0"/>
      <extension>
        <plugin plugin="mujoco.elasticity.cable"/>
      </extension>
      <worldbody>
        <geom type="plane" size="0 0 1" quat="1 0 0 0"/>
        <site name="reference" pos="0 0 0"/>
        <composite type="cable" curve="s" count="41 1 1" size="1" offset="0 0 1" initial="none">
          <plugin plugin="mujoco.elasticity.cable">
            <config key="twist" value="1e6"/>
            <config key="bend" value="1e9"/>
          </plugin>
          <joint kind="main" damping="2"/>
          <geom type="capsule" size=".005" density="1"/>
        </composite>
      </worldbody>
      <contact>
        <exclude body1="B_first" body2="B_last"/>
      </contact>
      <sensor>
        <framepos objtype="site" objname="S_last"/>
      </sensor>
      <actuator>
        <motor site="S_last" gear="0 0 0 0 1 0" ctrllimited="true" ctrlrange="0 4"/>
      </actuator>
    </mujoco>
    """
    physics = engine.Physics.from_xml_string(model_with_cable_plugin)
    physics_state = physics.get_state()

    # qpos, qvel, act, plugin_state
    self.assertLen(physics._physics_state_items(), 4)

    physics.set_state(physics_state)

    new_physics_state = np.random.random_sample(physics_state.shape)
    physics.set_state(new_physics_state)

    np.testing.assert_allclose(new_physics_state, physics.get_state())

  def testSetInvalidPhysicsState(self):
    badly_shaped_state = np.repeat(self._physics.get_state(), repeats=2)

    with self.assertRaises(ValueError):
      self._physics.set_state(badly_shaped_state)

  def testNamedIndexing(self):
    self.assertEqual((3,), self._physics.named.data.xpos['cart'].shape)
    self.assertEqual((2, 3),
                     self._physics.named.data.xpos[['cart', 'pole']].shape)

  def testReload(self):
    self._physics.reload_from_xml_path(MODEL_PATH)

  def testReset(self):
    self._physics.reset()
    self.assertEqual(self._physics.data.qpos[1], 0)
    keyframe_id = 0
    self._physics.reset(keyframe_id=keyframe_id)
    self.assertEqual(self._physics.data.qpos[1],
                     self._physics.model.key_qpos[keyframe_id, 1])
    out_of_range = [-1, 3]
    max_valid = self._physics.model.nkey - 1
    for actual in out_of_range:
      with self.assertRaisesWithLiteralMatch(
          ValueError,
          engine._KEYFRAME_ID_OUT_OF_RANGE.format(
              max_valid=max_valid, actual=actual)):
        self._physics.reset(keyframe_id=actual)

  def testLoadAndReloadFromStringWithAssets(self):
    physics = engine.Physics.from_xml_string(
        MODEL_WITH_ASSETS, assets=ASSETS)
    physics.reload_from_xml_string(MODEL_WITH_ASSETS, assets=ASSETS)

  @parameterized.parameters(*enums.mjtWarning._fields[:-1])
  def testDivergenceException(self, warning_name):
    warning_enum = getattr(enums.mjtWarning, warning_name)
    with self.assertRaisesWithLiteralMatch(
        control.PhysicsError,
        engine._INVALID_PHYSICS_STATE.format(warning_names=warning_name)):
      with self._physics.check_invalid_state():
        self._physics.data.warning[warning_enum].number = 1
    # Existing warnings should not raise an exception.
    with self._physics.check_invalid_state():
      pass
    self._physics.reset()
    with self._physics.check_invalid_state():
      pass

  @parameterized.parameters(float('inf'), float('nan'), 1e15)
  def testBadQpos(self, bad_value):
    with self._physics.reset_context():
      self._physics.data.qpos[0] = bad_value
    with self.assertRaises(control.PhysicsError):
      with self._physics.check_invalid_state():
        mujoco.mj_checkPos(self._physics.model.ptr, self._physics.data.ptr)
    self._physics.reset()
    with self._physics.check_invalid_state():
      mujoco.mj_checkPos(self._physics.model.ptr, self._physics.data.ptr)

  def testNanControl(self):
    with self._physics.reset_context():
      pass

    # Apply the controls.
    with self.assertRaisesWithLiteralMatch(
        control.PhysicsError,
        engine._INVALID_PHYSICS_STATE.format(warning_names='mjWARN_BADCTRL')):
      with self._physics.check_invalid_state():
        self._physics.data.ctrl[0] = float('nan')
        self._physics.step()

  def testSuppressPhysicsError(self):
    bad_value = float('nan')
    message = engine._INVALID_PHYSICS_STATE.format(
        warning_names='mjWARN_BADCTRL')

    def assert_physics_error():
      self._physics.data.ctrl[0] = bad_value
      with self.assertRaisesWithLiteralMatch(control.PhysicsError, message):
        self._physics.forward()

    def assert_warning():
      self._physics.data.ctrl[0] = bad_value
      with mock.patch.object(engine.logging, 'warn') as mock_warn:
        self._physics.forward()
      mock_warn.assert_called_once_with(message)

    assert_physics_error()
    with self._physics.suppress_physics_errors():
      assert_warning()
      with self._physics.suppress_physics_errors():
        assert_warning()
      assert_warning()
    assert_physics_error()

  @parameterized.named_parameters(
      ('_copy', lambda x: x.copy()),
      ('_deepcopy', copy.deepcopy),
      ('_pickle_and_unpickle', lambda x: pickle.loads(pickle.dumps(x))),
  )
  def testCopyOrPicklePhysics(self, func):
    for _ in range(10):
      self._physics.step()
    physics2 = func(self._physics)
    self.assertNotEqual(physics2.model.ptr, self._physics.model.ptr)
    self.assertNotEqual(physics2.data.ptr, self._physics.data.ptr)
    model_attr_to_compare = ('nnames', 'njmax', 'body_pos', 'geom_quat')
    self._assert_attributes_equal(
        physics2.model, self._physics.model, model_attr_to_compare)
    data_attr_to_compare = ('time', 'energy', 'qpos', 'xpos')
    self._assert_attributes_equal(
        physics2.data, self._physics.data, data_attr_to_compare)
    for _ in range(10):
      self._physics.step()
      physics2.step()
    self._assert_attributes_equal(
        physics2.model, self._physics.model, model_attr_to_compare)
    self._assert_attributes_equal(
        physics2.data, self._physics.data, data_attr_to_compare)

  @parameterized.named_parameters(
      ('_copy', lambda x: x.copy()),
      ('_pickle_and_unpickle', lambda x: pickle.loads(pickle.dumps(x))),
  )
  def testSuppressErrorsAfterCopyOrPicklePhysics(self, func):
    # Regression test for a problem that used to exist where
    # suppress_physics_errors couldn't be used on Physics objects that were
    # unpickled.
    physics2 = func(self._physics)
    with physics2.suppress_physics_errors():
      pass

  def testCopyDataOnly(self):
    physics2 = self._physics.copy(share_model=True)
    self.assertEqual(physics2.model.ptr, self._physics.model.ptr)
    self.assertNotEqual(physics2.data.ptr, self._physics.data.ptr)

  def testForwardDynamicsUpdatedAfterReset(self):
    gravity = -9.81
    self._physics.model.opt.gravity[2] = gravity
    with self._physics.reset_context():
      pass
    self.assertAlmostEqual(
        self._physics.named.data.sensordata['accelerometer'][2], -gravity)

  def testActuationNotAppliedInAfterReset(self):
    self._physics.data.ctrl[0] = 1.
    self._physics.after_reset()  # Calls `forward()` with actuation disabled.
    self.assertEqual(self._physics.data.actuator_force[0], 0.)
    self._physics.forward()  # Call `forward` directly with actuation enabled.
    self.assertEqual(self._physics.data.actuator_force[0], 1.)

  def testActionSpec(self):
    xml = """
    <mujoco>
      <worldbody>
        <body>
          <geom type="sphere" size="0.1"/>
          <joint type="hinge" name="hinge"/>
        </body>
      </worldbody>
      <actuator>
        <motor joint="hinge" ctrllimited="false"/>
        <motor joint="hinge" ctrllimited="true" ctrlrange="-1 2"/>
      </actuator>
    </mujoco>
    """
    physics = engine.Physics.from_xml_string(xml)
    spec = engine.action_spec(physics)
    self.assertEqual(float, spec.dtype)
    np.testing.assert_array_equal(spec.minimum, [-mujoco.mjMAXVAL, -1.0])
    np.testing.assert_array_equal(spec.maximum, [mujoco.mjMAXVAL, 2.0])

  def testNstep(self):
    # Make initial state.
    with self._physics.reset_context():
      self._physics.data.qvel[0] = 1
      self._physics.data.qvel[1] = 1
    initial_state = self._physics.get_state()

    # step() 4 times.
    for _ in range(4):
      self._physics.step()
    for_loop_state = self._physics.get_state()

    # Reset state, call step(4).
    with self._physics.reset_context():
      self._physics.set_state(initial_state)
    self._physics.step(4)
    nstep_state = self._physics.get_state()

    np.testing.assert_array_equal(for_loop_state, nstep_state)

    # Repeat test with with RK4 integrator:
    self._physics.model.opt.integrator = enums.mjtIntegrator.mjINT_RK4

    # step() 4 times.
    with self._physics.reset_context():
      self._physics.set_state(initial_state)
    for _ in range(4):
      self._physics.step()
    for_loop_state_rk4 = self._physics.get_state()

    # Reset state, call step(4).
    with self._physics.reset_context():
      self._physics.set_state(initial_state)
    self._physics.step(4)
    nstep_state_rk4 = self._physics.get_state()

    np.testing.assert_array_equal(for_loop_state_rk4, nstep_state_rk4)

if __name__ == '__main__':
  absltest.main()