mujoco

# Copyright 2022 DeepMind Technologies Limited
#
# 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 the MuJoCo renderer."""

import gc
import sys

from absl.testing import absltest
from absl.testing import parameterized
import mujoco
import numpy as np


@absltest.skipUnless(
    hasattr(mujoco, 'GLContext'), 'MuJoCo rendering is disabled'
)
class MuJoCoRendererTest(parameterized.TestCase):

  def test_renderer_info_binding(self):
    info = mujoco.MjrRendererInfo()
    mujoco.mjr_getRendererInfo(info)

    self.assertIn(info.renderer, ('classic', 'filament', 'noop'))
    self.assertIn(info.backend, ('', 'opengl', 'vulkan', 'unknown'))

  def test_renderer_unknown_camera_name(self):
    xml = """
<mujoco>
  <worldbody>
    <camera name="a"/>
  </worldbody>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    with mujoco.Renderer(model, 50, 50) as renderer:
      mujoco.mj_forward(model, data)
      with self.assertRaisesRegex(ValueError, r'camera "b" does not exist'):
        renderer.update_scene(data, 'b')

  def test_renderer_camera_under_range(self):
    xml = """
<mujoco>
  <worldbody>
    <camera name="a"/>
  </worldbody>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    with mujoco.Renderer(model, 50, 50) as renderer:
      mujoco.mj_forward(model, data)
      with self.assertRaisesRegex(ValueError, '-2 is out of range'):
        renderer.update_scene(data, -2)

  def test_renderer_camera_over_range(self):
    xml = """
<mujoco>
  <worldbody>
    <camera name="a"/>
  </worldbody>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    with mujoco.Renderer(model, 50, 50) as renderer:
      mujoco.mj_forward(model, data)
      with self.assertRaisesRegex(ValueError, '1 is out of range'):
        renderer.update_scene(data, 1)

  def test_renderer_renders_scene(self):
    xml = """
<mujoco>
  <worldbody>
    <camera name="closeup" pos="0 -6 0" xyaxes="1 0 0 0 1 100"/>
    <geom name="white_box" type="box" size="1 1 1" rgba="1 1 1 1"/>
  </worldbody>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    with mujoco.Renderer(model, 50, 50) as renderer:
      mujoco.mj_forward(model, data)
      renderer.update_scene(data, 'closeup')

      pixels = renderer.render().flatten()
      not_all_black = False

      # Pixels should all be a neutral color.
      for pixel in pixels:
        if pixel > 0:
          not_all_black = True
          break
      self.assertTrue(not_all_black)

  def test_renderer_output_without_out(self):
    xml = """
<mujoco>
  <worldbody>
    <camera name="closeup" pos="0 -6 0" xyaxes="1 0 0 0 1 100"/>
    <geom name="white_box" type="box" size="1 1 1" rgba="1 1 1 1"/>
  </worldbody>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    mujoco.mj_forward(model, data)
    with mujoco.Renderer(model, 50, 50) as renderer:
      renderer.update_scene(data, 'closeup')
      pixels = [renderer.render()]

      colors = (
          (1.0, 0.0, 0.0, 1.0),
          (0.0, 1.0, 0.0, 1.0),
          (0.0, 0.0, 1.0, 1.0),
      )

      for i, color in enumerate(colors):
        model.geom_rgba[0, :] = color
        mujoco.mj_forward(model, data)
        renderer.update_scene(data, 'closeup')
        pixels.append(renderer.render())
        self.assertIsNot(pixels[-2], pixels[-1])

        # Pixels should change over steps.
        self.assertFalse((pixels[i + 1] == pixels[i]).all())

  def test_renderer_output_with_out(self):
    xml = """
<mujoco>
  <worldbody>
    <camera name="closeup" pos="0 -6 0" xyaxes="1 0 0 0 1 100"/>
    <geom name="white_box" type="box" size="1 1 1" rgba="1 1 1 1"/>
  </worldbody>
</mujoco>
"""
    render_size = (50, 50)
    render_out = np.zeros((*render_size, 3), np.uint8)
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    mujoco.mj_forward(model, data)
    with mujoco.Renderer(model, *render_size) as renderer:
      renderer.update_scene(data, 'closeup')

      self.assertTrue(np.all(render_out == 0))

      pixels = renderer.render(out=render_out)

      # Pixels should always refer to the same `render_out` array.
      self.assertIs(pixels, render_out)
      self.assertFalse(np.all(render_out == 0))

      failing_render_size = (10, 10)
      self.assertNotEqual(failing_render_size, render_size)
      with self.assertRaises(ValueError):
        renderer.render(out=np.zeros((*failing_render_size, 3), np.uint8))

  def test_renderer_depth_rendering_camera_projection(self):
    """Tests orthographic vs perspective camera depth rendering projections.

    Orthographic depth uses a linear NDC mapping (glOrtho); perspective depth
    uses a hyperbolic mapping (glFrustum).

    Render the same tilted plane using identically positioned orthographic and
    perspective cameras. The two should agree on the depth at the center pixel
    (actual distance), but they should have different depth values at non-center
    pixels because the rays of the perspective camera diverge while those of the
    orthographic camera remain parallel.
    """
    distance = 3.0
    ortho_fovy = 2.0 * distance * np.tan(np.deg2rad(45.0 / 2.0))
    xml = f"""
<mujoco>
  <statistic extent="1"/>
  <visual>
    <quality offsamples="0"/>
  </visual>
  <worldbody>
    <camera name="ortho" pos="0 0 {distance}" xyaxes="1 0 0 0 1 0" projection="orthographic" fovy="{ortho_fovy}"/>
    <camera name="persp" pos="0 0 {distance}" xyaxes="1 0 0 0 1 0" projection="perspective" fovy="45"/>
    <geom name="floor" type="plane" size="5 5 0.1" zaxis="0 0.3 1"/>
  </worldbody>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    mujoco.mj_forward(model, data)

    # Odd resolution so a single pixel sits exactly on the optical axis.
    def render_depth(camera_name):
      with mujoco.Renderer(model, 51, 51) as renderer:
        renderer.enable_depth_rendering()
        renderer.update_scene(data, camera_name)
        return renderer.scene.camera[0].orthographic, renderer.render().copy()

    is_ortho, ortho_depth = render_depth('ortho')
    is_persp, persp_depth = render_depth('persp')

    # Check that the scenes are actually rendered with the expected cameras.
    self.assertTrue(is_ortho)
    self.assertFalse(is_persp)

    # Center pixels of both cameras should show the actual distance
    center = ortho_depth.shape[0] // 2, ortho_depth.shape[1] // 2
    self.assertAlmostEqual(ortho_depth[center], distance, delta=1e-2)
    self.assertAlmostEqual(persp_depth[center], distance, delta=1e-2)

    # Depth values at the corner should differ
    corner = (0, 0)
    self.assertNotAlmostEqual(
        ortho_depth[corner], persp_depth[corner], delta=1e-2
    )

  def test_renderer_del_safe_when_init_fails_early(self):
    """Regression test for #3213.

    Renderer.__del__ must be safe on a partially-constructed instance when
    __init__ raises before the rendering contexts are assigned. Previously,
    AttributeError from __del__ masked the real __init__ exception.
    """
    xml = """
<mujoco>
  <visual>
    <global offwidth="50" offheight="50"/>
  </visual>
  <worldbody/>
</mujoco>
"""
    model = mujoco.MjModel.from_xml_string(xml)

    # Capture any exception raised from __del__ on the partially-constructed
    # Renderer. Without the fix, __del__ raises AttributeError, which is
    # funneled through sys.unraisablehook.
    unraisable = []
    old_hook = sys.unraisablehook
    sys.unraisablehook = lambda args: unraisable.append(args)
    try:
      # width > offwidth raises ValueError in __init__ before
      # self._gl_context is assigned.
      with self.assertRaises(ValueError):
        mujoco.Renderer(model, height=50, width=200)
      gc.collect()
    finally:
      sys.unraisablehook = old_hook

    self.assertEqual(
        [u.exc_type.__name__ for u in unraisable],
        [],
        msg=(
            'Renderer.__del__ raised on a partially-constructed instance; '
            'see #3213.'
        ),
    )


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