Merge pull request #978 from aftersomemath:depth-precision-in-place

PiperOrigin-RevId: 572961667
Change-Id: I69011ddf835ae286e30e7f84cca67867cfe94fa4

python/mujoco/render.cc

@@ -215,6 +215,7 @@ PYBIND11_MODULE(_render, pymodule) {
   X(windowDoublebuffer);
   X(currentBuffer);
   X(readPixelFormat);
+  X(readDepthMap);
 #undef X
 
 #define X(var)                      \

python/mujoco/renderer.py

@@ -86,6 +86,7 @@ the clause:
     _render.mjr_setBuffer(
         _enums.mjtFramebuffer.mjFB_OFFSCREEN.value, self._mjr_context
     )
+    self._mjr_context.readDepthMap = _enums.mjtDepthMap.mjDEPTH_ZEROFAR
 
     # Default render flags.
     self._depth_rendering = False
@@ -138,7 +139,9 @@ the clause:
     """
     original_flags = self._scene.flags.copy()
 
-    if self._segmentation_rendering:
+    # Using segmented rendering for depth makes the calculated depth more
+    # accurate at far distances.
+    if self._depth_rendering or self._segmentation_rendering:
       self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True
       self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True
 
@@ -173,11 +176,34 @@ the clause:
       near = self._model.vis.map.znear * extent
       far = self._model.vis.map.zfar * extent
 
-      # Convert from [0 1] to depth in units of length, see links below:
-      # http://stackoverflow.com/a/6657284/1461210
-      # https://www.khronos.org/opengl/wiki/Depth_Buffer_Precision
-      out = near / (1 - out * (1 - near / far))
+      # Calculate OpenGL perspective matrix values in float32 precision
+      # so they are close to what glFrustum returns
+      # https://registry.khronos.org/OpenGL-Refpages/gl2.1/xhtml/glFrustum.xml
+      zfar = np.float32(far)
+      znear = np.float32(near)
+      c_coef = -(zfar + znear) / (zfar - znear)
+      d_coef = -(np.float32(2) * zfar * znear) / (zfar - znear)
 
+      # In reverse Z mode the perspective matrix is transformed by the following
+      c_coef = np.float32(-0.5) * c_coef - np.float32(0.5)
+      d_coef = np.float32(-0.5) * d_coef
+
+      # We need 64 bits to convert Z from ndc to metric depth without noticeable
+      # losses in precision
+      out_64 = out.astype(np.float64)
+
+      # Undo OpenGL projection
+      # Note: We do not need to take action to convert from window coordinates
+      # to normalized device coordinates because in reversed Z mode the mapping
+      # is identity
+      out_64 = d_coef / (out_64 + c_coef)
+
+      # Cast result back to float32 for backwards compatibility
+      # This has a small accuracy cost
+      out[:] = out_64.astype(np.float32)
+
+      # Reset scene flags.
+      np.copyto(self._scene.flags, original_flags)
     elif self._segmentation_rendering:
       _render.mjr_readPixels(out, None, self._rect, self._mjr_context)