# Copyright 2018 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.mjcf.physics`."""

import copy
import os
import pickle

from absl.testing import absltest
from absl.testing import parameterized
from dm_control import mjcf
from dm_control.mjcf import physics as mjcf_physics
from dm_control.mujoco.wrapper import mjbindings
import mock
import mujoco
import numpy as np

mjlib = mjbindings.mjlib

ARM_MODEL = os.path.join(os.path.dirname(__file__), 'test_assets/robot_arm.xml')


class PhysicsTest(parameterized.TestCase):
  """Tests for `mjcf.Physics`."""

  def setUp(self):
    super().setUp()
    self.model = mjcf.from_path(ARM_MODEL)
    self.physics = mjcf.Physics.from_xml_string(
        self.model.to_xml_string(), assets=self.model.get_assets())
    self.random = np.random.RandomState(0)

  def sample_elements(self, namespace, single_element):
    all_elements = self.model.find_all(namespace)
    if single_element:
      # A single randomly chosen element from this namespace.
      elements = self.random.choice(all_elements)
      full_identifiers = elements.full_identifier
    else:
      # A random permutation of all elements in this namespace.
      elements = self.random.permutation(all_elements)
      full_identifiers = [element.full_identifier for element in elements]
    return elements, full_identifiers

  def test_construct_and_reload_from_mjcf_model(self):
    physics = mjcf.Physics.from_mjcf_model(self.model)
    physics.data.time = 1.
    physics.reload_from_mjcf_model(self.model)
    self.assertEqual(physics.data.time, 0.)

  @parameterized.parameters(
      # namespace, single_element
      ('geom', True),
      ('geom', False),
      ('joint', True),
      ('joint', False),
      ('actuator', True),
      ('actuator', False))
  def test_id(self, namespace, single_element):
    elements, full_identifiers = self.sample_elements(namespace, single_element)
    actual = self.physics.bind(elements).element_id
    if single_element:
      expected = self.physics.model.name2id(full_identifiers, namespace)
    else:
      expected = [self.physics.model.name2id(name, namespace)
                  for name in full_identifiers]
    np.testing.assert_array_equal(expected, actual)

  def assertCanGetAndSetBindingArray(
      self, binding, attribute_name, named_indexer, full_identifiers):

    # Read the values using the binding attribute.
    actual = getattr(binding, attribute_name)

    # Read them using the normal named indexing machinery.
    expected = named_indexer[full_identifiers]
    np.testing.assert_array_equal(expected, actual)

    # Assign an array of unique values to the attribute.
    expected = np.arange(actual.size, dtype=actual.dtype).reshape(actual.shape)
    setattr(binding, attribute_name, expected)

    # Read the values back using the normal named indexing machinery.
    actual = named_indexer[full_identifiers]
    np.testing.assert_array_equal(expected, actual)

  @parameterized.parameters(
      # namespace, attribute_name, model_or_data, field_name, single_element
      ('geom', 'xpos', 'data', 'geom_xpos', True),
      ('geom', 'xpos', 'data', 'geom_xpos', False),
      ('joint', 'qpos', 'data', 'qpos', True),
      ('joint', 'qpos', 'data', 'qpos', False),
      ('site', 'rgba', 'model', 'site_rgba', True),
      ('site', 'rgba', 'model', 'site_rgba', False),
      ('sensor', 'sensordata', 'data', 'sensordata', True),
      ('sensor', 'sensordata', 'data', 'sensordata', False))
  def test_attribute_access(self, namespace, attribute_name, model_or_data,
                            field_name, single_element):

    elements, full_identifiers = self.sample_elements(namespace, single_element)
    named_indexer = getattr(getattr(self.physics.named, model_or_data),
                            field_name)
    binding = self.physics.bind(elements)
    self.assertCanGetAndSetBindingArray(
        binding, attribute_name, named_indexer, full_identifiers)

  @parameterized.parameters(
      # namespace, attribute_name, model_or_data, field_name, single_element,
      # column_index
      ('geom', 'pos', 'model', 'geom_pos', True, None),
      ('geom', 'pos', 'model', 'geom_pos', False, None),
      ('geom', 'pos', 'model', 'geom_pos', True, 1),
      ('geom', 'pos', 'model', 'geom_pos', False, 1),
      ('geom', 'pos', 'model', 'geom_pos', True, 'y'),
      ('geom', 'pos', 'model', 'geom_pos', False, 'y'),
      ('geom', 'pos', 'model', 'geom_pos', True, slice(0, None, 2)),
      ('geom', 'pos', 'model', 'geom_pos', False, slice(0, None, 2)),
      ('geom', 'pos', 'model', 'geom_pos', True, [0, 2]),
      ('geom', 'pos', 'model', 'geom_pos', False, [0, 2]),
      ('geom', 'pos', 'model', 'geom_pos', True, ['x', 'z']),
      ('geom', 'pos', 'model', 'geom_pos', False, ['x', 'z']),
      ('joint', 'qpos', 'data', 'qpos', True, None),
      ('joint', 'qpos', 'data', 'qpos', False, None))
  def test_indexing(self, namespace, attribute_name, model_or_data,
                    field_name, single_element, column_index):

    elements, full_identifiers = self.sample_elements(namespace, single_element)
    named_indexer = getattr(getattr(self.physics.named, model_or_data),
                            field_name)
    binding = self.physics.bind(elements)

    if column_index is not None:
      binding_index = (attribute_name, column_index)
      try:
        named_index = np.ix_(full_identifiers, column_index)
      except ValueError:
        named_index = (full_identifiers, column_index)
    else:
      binding_index = attribute_name
      named_index = full_identifiers

    # Read the values by indexing the binding.
    actual = binding[binding_index]

    # Read them using the normal named indexing machinery.
    expected = named_indexer[named_index]
    np.testing.assert_array_equal(expected, actual)

    # Write an array of unique values into the binding.
    expected = np.arange(actual.size, dtype=actual.dtype).reshape(actual.shape)
    binding[binding_index] = expected

    # Read the values back using the normal named indexing machinery.
    actual = named_indexer[named_index]
    np.testing.assert_array_equal(expected, actual)

  def test_bind_mocap_body(self):
    pos = [1, 2, 3]
    quat = [1, 0, 0, 0]
    model = mjcf.RootElement()
    # Bodies are non-mocap by default.
    normal_body = model.worldbody.add('body', pos=pos, quat=quat)
    mocap_body = model.worldbody.add('body', pos=pos, quat=quat, mocap='true')
    physics = mjcf.Physics.from_xml_string(model.to_xml_string())

    binding = physics.bind(mocap_body)
    np.testing.assert_array_equal(pos, binding.mocap_pos)
    np.testing.assert_array_equal(quat, binding.mocap_quat)

    new_pos = [4, 5, 6]
    new_quat = [0, 1, 0, 0]
    binding.mocap_pos = new_pos
    binding.mocap_quat = new_quat
    np.testing.assert_array_equal(
        new_pos, physics.named.data.mocap_pos[mocap_body.full_identifier])
    np.testing.assert_array_equal(
        new_quat, physics.named.data.mocap_quat[mocap_body.full_identifier])

    with self.assertRaises(AttributeError):
      _ = physics.bind(normal_body).mocap_pos

    with self.assertRaisesRegex(
        ValueError,
        'Cannot bind to a collection containing multiple element types'):
      physics.bind([mocap_body, normal_body])

  def test_bind_worldbody(self):
    expected_mass = 10
    model = mjcf.RootElement()
    child = model.worldbody.add('body')
    child.add('geom', type='sphere', size=[0.1], mass=expected_mass)
    physics = mjcf.Physics.from_mjcf_model(model)
    mass = physics.bind(model.worldbody).subtreemass
    self.assertEqual(mass, expected_mass)

  def test_bind_stateful_actuator(self):
    model = mjcf.RootElement()
    body = model.worldbody.add('body')
    body.add('joint', name='joint')
    body.add('geom', type='sphere', size=[1])

    model.actuator.add(
        'general', name='act1', joint='joint', dyntype='integrator')

    physics = mjcf.Physics.from_mjcf_model(model)
    actuator = model.find('actuator', 'act1')
    binding = physics.bind(actuator)

    # This used to crash
    self.assertEqual(0, binding.act)

  def test_caching(self):
    all_joints = self.model.find_all('joint')

    original = self.physics.bind(all_joints)
    cached = self.physics.bind(all_joints)
    self.assertIs(cached, original)

    different_order = self.physics.bind(all_joints[::-1])
    self.assertIsNot(different_order, original)

    # Reloading the `Physics` instance should clear the cache.
    self.physics.reload_from_xml_string(
        self.model.to_xml_string(), assets=self.model.get_assets())
    after_reload = self.physics.bind(all_joints)
    self.assertIsNot(after_reload, original)

  def test_exceptions(self):
    joint = self.model.find_all('joint')[0]
    geom = self.model.find_all('geom')[0]
    with self.assertRaisesRegex(
        ValueError,
        'Cannot bind to a collection containing multiple element types'):
      self.physics.bind([joint, geom])

    with self.assertRaisesRegex(ValueError, 'cannot be bound to physics'):
      mjcf.physics.Binding(self.physics, 'invalid_namespace', 'whatever')

    binding = self.physics.bind(joint)
    with self.assertRaisesRegex(AttributeError, 'does not have attribute'):
      getattr(binding, 'invalid_attribute')

  def test_dirty(self):
    self.physics.forward()
    self.assertFalse(self.physics.is_dirty)

    joints, _ = self.sample_elements('joint', single_element=False)
    sites, _ = self.sample_elements('site', single_element=False)

    # Accessing qpos shouldn't trigger a recalculation.
    _ = self.physics.bind(joints).qpos
    self.assertFalse(self.physics.is_dirty)

    # Reassignments to qpos should cause the physics to become dirty.
    site_xpos_before = copy.deepcopy(self.physics.bind(sites).xpos)
    self.physics.bind(joints).qpos += 0.5
    self.assertTrue(self.physics.is_dirty)

    # Accessing stuff in mjModel shouldn't trigger a recalculation.
    _ = self.physics.bind(sites).pos
    self.assertTrue(self.physics.is_dirty)

    # Accessing stuff in mjData should trigger a recalculation.
    actual_sites_xpos_after = copy.deepcopy(self.physics.bind(sites).xpos)
    self.assertFalse(self.physics.is_dirty)
    self.assertFalse((actual_sites_xpos_after == site_xpos_before).all())

    # Automatic recalculation should render `forward` a no-op here.
    self.physics.forward()
    expected_sites_xpos_after = self.physics.bind(sites).xpos
    np.testing.assert_array_equal(actual_sites_xpos_after,
                                  expected_sites_xpos_after)

    # `forward` should not be called on subsequent queries to xpos.
    with mock.patch.object(
        self.physics, 'forward',
        side_effect=self.physics.forward) as mock_forward:
      _ = self.physics.bind(sites).xpos
      mock_forward.assert_not_called()

  @parameterized.parameters(True, False)
  def test_assign_while_dirty(self, assign_via_slice):
    actuators = self.model.find_all('actuator')
    if assign_via_slice:
      self.physics.bind(actuators).ctrl[:] = 0.75
    else:
      self.physics.bind(actuators).ctrl = 0.75
    self.assertTrue(self.physics.is_dirty)
    self.physics.step()
    self.assertTrue(self.physics.is_dirty)
    sensors = self.model.find_all('sensor')
    if assign_via_slice:
      self.physics.bind(sensors).sensordata[:] = 12345
    else:
      self.physics.bind(sensors).sensordata = 12345
    self.assertFalse(self.physics.is_dirty)
    np.testing.assert_array_equal(
        self.physics.bind(sensors).sensordata,
        [12345] * len(self.physics.bind(sensors).sensordata))

  def test_setitem_on_binding_attr(self):
    bodies, _ = self.sample_elements('body', single_element=False)
    xfrc_binding = self.physics.bind(bodies).xfrc_applied

    xfrc_binding[:, 1] = list(range(len(bodies)))
    for i, body in enumerate(bodies):
      self.assertEqual(xfrc_binding[i, 1], i)
      self.assertEqual(
          self.physics.named.data.xfrc_applied[body.full_identifier][1], i)

    xfrc_binding[:, 1] *= 2
    for i, body in enumerate(bodies):
      self.assertEqual(xfrc_binding[i, 1], 2 * i)
      self.assertEqual(
          self.physics.named.data.xfrc_applied[body.full_identifier][1], 2 * i)

    xfrc_binding[[1, 3, 5], 2] = 42
    for i, body in enumerate(bodies):
      actual_value = (
          self.physics.named.data.xfrc_applied[body.full_identifier][2])
      if i in [1, 3, 5]:
        self.assertEqual(actual_value, 42)
      else:
        self.assertNotEqual(actual_value, 42)

    # Bind to a single element.
    single_binding = self.physics.bind(bodies[0]).xfrc_applied
    single_binding[:2] = 55
    np.testing.assert_array_equal(single_binding[:2], [55, 55])
    np.testing.assert_array_equal(
        self.physics.named.data.xfrc_applied[bodies[0].full_identifier][:2],
        [55, 55])

  def test_empty_binding(self):
    binding = self.physics.bind([])
    self.assertEqual(binding.xpos.shape, (0,))
    with self.assertRaisesWithLiteralMatch(
        ValueError, 'Cannot assign a value to an empty binding.'):
      binding.xpos = 5

  @parameterized.parameters([('data', 'act'), ('data', 'act_dot')])
  def test_actuator_state_binding(self, model_or_data, attribute_name):

    def make_model_with_mixed_actuators():
      actuators = []
      is_stateful = []
      root = mjcf.RootElement()
      body = root.worldbody.add('body')
      body.add('geom', type='sphere', size=[0.1])
      slider = body.add('joint', type='slide', name='slide_joint')
      # Third-order `general` actuator.
      actuators.append(
          root.actuator.add(
              'general', dyntype='integrator', biastype='affine',
              dynprm=[1, 0, 0], joint=slider, name='general_act'))
      is_stateful.append(True)
      # Cylinder actuators are also third-order.
      actuators.append(
          root.actuator.add('cylinder', joint=slider, name='cylinder_act'))
      is_stateful.append(True)
      # A second-order actuator, added after the third-order actuators. The
      # actuators will be automatically reordered in the generated XML so that
      # the second-order actuator comes first.
      actuators.append(
          root.actuator.add('velocity', joint=slider, name='velocity_act'))
      is_stateful.append(False)
      return root, actuators, is_stateful

    model, actuators, is_stateful = make_model_with_mixed_actuators()
    physics = mjcf.Physics.from_mjcf_model(model)
    binding = physics.bind(actuators)
    named_indexer = getattr(
        getattr(physics.named, model_or_data), attribute_name)
    stateful_actuator_names = [
        actuator.full_identifier
        for actuator, stateful in zip(actuators, is_stateful) if stateful]
    self.assertCanGetAndSetBindingArray(
        binding, attribute_name, named_indexer, stateful_actuator_names)

  def test_bind_stateless_actuators_only(self):
    actuators = []
    root = mjcf.RootElement()
    body = root.worldbody.add('body')
    body.add('geom', type='sphere', size=[0.1])
    slider = body.add('joint', type='slide', name='slide_joint')
    actuators.append(
        root.actuator.add('velocity', joint=slider, name='velocity_act'))
    actuators.append(
        root.actuator.add('motor', joint=slider, name='motor_act'))
    # `act` should be an empty array if there are no stateful actuators.
    physics = mjcf.Physics.from_mjcf_model(root)
    self.assertEqual(physics.bind(actuators).act.shape, (0,))

  def make_simple_model(self):
    def add_submodel(root):
      body = root.worldbody.add('body')
      geom = body.add('geom', type='ellipsoid', size=[0.1, 0.2, 0.3])
      site = body.add('site', type='sphere', size=[0.1])
      return body, geom, site
    root = mjcf.RootElement()
    add_submodel(root)
    add_submodel(root)
    return root

  def quat2mat(self, quat):
    result = np.empty(9, dtype=np.double)
    mujoco.mju_quat2Mat(result, np.asarray(quat))
    return result

  @parameterized.parameters(['body', 'geom', 'site'])
  def test_write_to_pos(self, entity_type):
    root = self.make_simple_model()
    entity1, entity2 = root.find_all(entity_type)
    physics = mjcf.Physics.from_mjcf_model(root)
    first = physics.bind(entity1)
    second = physics.bind(entity2)

    # Initially both entities should be 'sameframe'
    self.assertEqual(first.sameframe, 1)
    self.assertEqual(second.sameframe, 1)

    # Assigning to `pos` should disable 'sameframe' only for that entity.
    new_pos = (0., 0., 0.1)
    first.pos = new_pos
    self.assertEqual(first.sameframe, 0)
    self.assertEqual(second.sameframe, 1)
    # `xpos` should reflect the new position.
    np.testing.assert_array_equal(first.xpos, new_pos)

    # Writing into the `pos` array should also disable 'sameframe'.
    new_x = -0.1
    pos_array = second.pos
    pos_array[0] = new_x
    self.assertEqual(second.sameframe, 0)
    # `xpos` should reflect the new position.
    self.assertEqual(second.xpos[0], new_x)

  @parameterized.parameters(['body', 'geom', 'site'])
  def test_write_to_quat(self, entity_type):
    root = self.make_simple_model()
    entity1, entity2 = root.find_all(entity_type)
    physics = mjcf.Physics.from_mjcf_model(root)
    first = physics.bind(entity1)
    second = physics.bind(entity2)

    # Initially both entities should be 'sameframe'
    self.assertEqual(first.sameframe, 1)
    self.assertEqual(second.sameframe, 1)

    # Assigning to `quat` should disable 'sameframe' only for that entity.
    new_quat = (0., 0., 0., 1.)
    first.quat = new_quat
    self.assertEqual(first.sameframe, 0)
    self.assertEqual(second.sameframe, 1)
    # `xmat` should reflect the new quaternion.
    np.testing.assert_allclose(first.xmat, self.quat2mat(new_quat))

    # Writing into the `quat` array should also disable 'sameframe'.
    new_w = -1.
    quat_array = second.quat
    quat_array[0] = new_w
    self.assertEqual(second.sameframe, 0)
    # `xmat` should reflect the new quaternion.
    np.testing.assert_allclose(second.xmat, self.quat2mat(quat_array))

  def test_write_to_ipos(self):
    root = self.make_simple_model()
    entity1, entity2 = root.find_all('body')
    physics = mjcf.Physics.from_mjcf_model(root)
    first = physics.bind(entity1)
    second = physics.bind(entity2)

    # Initially both bodies should be 'simple' and 'sameframe'
    self.assertEqual(first.simple, 1)
    self.assertEqual(first.sameframe, 1)
    self.assertEqual(second.simple, 1)
    self.assertEqual(second.sameframe, 1)

    # Assigning to `ipos` should disable 'simple' and 'sameframe' only for that
    # body.
    new_ipos = (0., 0., 0.1)
    first.ipos = new_ipos
    self.assertEqual(first.simple, 0)
    self.assertEqual(first.sameframe, 0)
    self.assertEqual(second.simple, 1)
    self.assertEqual(second.sameframe, 1)
    # `xipos` should reflect the new position.
    np.testing.assert_array_equal(first.xipos, new_ipos)

    # Writing into the `ipos` array should also disable 'simple' and
    # 'sameframe'.
    new_x = -0.1
    ipos_array = second.ipos
    ipos_array[0] = new_x
    self.assertEqual(second.simple, 0)
    self.assertEqual(second.sameframe, 0)
    # `xipos` should reflect the new position.
    self.assertEqual(second.xipos[0], new_x)

  def test_write_to_iquat(self):
    root = self.make_simple_model()
    entity1, entity2 = root.find_all('body')
    physics = mjcf.Physics.from_mjcf_model(root)
    first = physics.bind(entity1)
    second = physics.bind(entity2)

    # Initially both bodies should be 'simple' and 'sameframe'
    self.assertEqual(first.simple, 1)
    self.assertEqual(first.sameframe, 1)
    self.assertEqual(second.simple, 1)
    self.assertEqual(second.sameframe, 1)

    # Assigning to `iquat` should disable 'simple' and 'sameframe' only for that
    # body.
    new_iquat = (0., 0., 0., 1.)
    first.iquat = new_iquat
    self.assertEqual(first.simple, 0)
    self.assertEqual(first.sameframe, 0)
    self.assertEqual(second.simple, 1)
    self.assertEqual(second.sameframe, 1)
    # `ximat` should reflect the new quaternion.
    np.testing.assert_allclose(first.ximat, self.quat2mat(new_iquat))

    # Writing into the `iquat` array should also disable 'simple' and
    # 'sameframe'.
    new_w = -0.1
    iquat_array = second.iquat
    iquat_array[0] = new_w
    self.assertEqual(second.simple, 0)
    self.assertEqual(second.sameframe, 0)
    # `ximat` should reflect the new quaternion.
    np.testing.assert_allclose(second.ximat, self.quat2mat(iquat_array))

  @parameterized.parameters([dict(order='C'), dict(order='F')])
  def test_copy_synchronizing_array_wrapper(self, order):
    root = self.make_simple_model()
    physics = mjcf.Physics.from_mjcf_model(root)
    xpos_view = physics.bind(root.find_all('body')).xpos
    xpos_view_copy = xpos_view.copy(order=order)

    np.testing.assert_array_equal(xpos_view, xpos_view_copy)
    self.assertFalse(np.may_share_memory(xpos_view, xpos_view_copy),
                     msg='Original and copy should not share memory.')
    self.assertIs(type(xpos_view_copy), np.ndarray)

    # Check that `order=` is respected.
    if order == 'C':
      self.assertTrue(xpos_view_copy.flags.c_contiguous)
      self.assertFalse(xpos_view_copy.flags.f_contiguous)
    elif order == 'F':
      self.assertFalse(xpos_view_copy.flags.c_contiguous)
      self.assertTrue(xpos_view_copy.flags.f_contiguous)

    # The copy should be writeable.
    self.assertTrue(xpos_view_copy.flags.writeable)
    new_value = 99.
    xpos_view_copy[0, -1] = new_value
    self.assertEqual(xpos_view_copy[0, -1], new_value)

  def test_error_when_pickling_synchronizing_array_wrapper(self):
    root = self.make_simple_model()
    physics = mjcf.Physics.from_mjcf_model(root)
    xpos_view = physics.bind(root.find_all('body')).xpos
    with self.assertRaisesWithLiteralMatch(
        NotImplementedError,
        mjcf_physics._PICKLING_NOT_SUPPORTED.format(type=type(xpos_view))):
      pickle.dumps(xpos_view)

  def test_plugins_elasticity(self):
    root = mjcf.RootElement()
    root.extension.add('plugin', plugin='mujoco.elasticity.cable')

    # Replicate example in mujoco/model/plugin/elasticity/cable.xml
    composite = root.worldbody.add(
        'composite',
        type='cable',
        curve='s',
        count=[41, 1, 1],
        size=[1, 0, 0],
        offset=[-0.3, 0, 0.6],
        initial='none',
    )
    plugin = composite.add('plugin', plugin='mujoco.elasticity.cable')
    plugin.add('config', key='twist', value='1e7')
    plugin.add('config', key='bend', value='4e6')
    plugin.add('config', key='vmax', value='0.05')
    composite.add('joint', kind='main', damping=0.015)
    composite.geom.type = 'capsule'
    composite.geom.size = [0.005, 0, 0]
    composite.geom.rgba = [0.8, 0.2, 0.1, 1]
    composite.geom.condim = 1

    physics = mjcf.Physics.from_mjcf_model(root)
    physics.step()

  def test_plugins_sdf(self):
    root = mjcf.RootElement()
    root.option.sdf_iterations = 10
    root.option.sdf_initpoints = 40

    extension = root.extension.add('plugin', plugin='mujoco.sdf.torus')
    instance = extension.add('instance', name='torus')
    instance.add('config', key='radius1', value='0.35')
    instance.add('config', key='radius2', value='0.15')

    # Replicate example in mujoco/model/plugin/elasticity/torus.xml
    mesh = root.asset.add('mesh', name='torus')
    mesh.add('plugin', instance='torus')

    # Test we can add SDF geom to the worldbody.
    worldbody_geom = root.worldbody.add(
        'geom', type='sdf', mesh='torus', rgba=[.2, .2, .8, 1])
    worldbody_geom.add('plugin', instance='torus')

    # Test we can add SDF geom to a body.
    body = root.worldbody.add('body', pos=[-1, 0, 3.8])
    body.add('freejoint')
    body_geom = body.add('geom', type='sdf', mesh='torus', rgba=[.2, .2, .8, 1])
    body_geom.add('plugin', instance='torus')

    physics = mjcf.Physics.from_mjcf_model(root)
    physics.step()

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