Fix extra dim for depth, matches rgb output shape.

PiperOrigin-RevId: 878098932
Change-Id: Ifcc4ba39dbee1beeea88fda44c3d20b88179705e

mjx/mujoco/mjx/_src/render_util.py

@@ -110,4 +110,4 @@ def get_depth(
   )
 
   depth = jnp.clip(raw / depth_scale, 0.0, 1.0)
-  return depth.reshape(height, width)
+  return depth.reshape(height, width, 1)

mjx/mujoco/mjx/_src/render_util_test.py

@@ -94,7 +94,7 @@ class RenderUtilTest(absltest.TestCase):
           rc, 0, depth_data, 5.0
       )
 
-    self.assertEqual(depth.shape, (height, width))
+    self.assertEqual(depth.shape, (height, width, 1))
 
   def test_get_depth_vmap(self):
     nworld, width, height = 3, 4, 4
@@ -111,7 +111,7 @@ class RenderUtilTest(absltest.TestCase):
           static_argnums=(0, 1, 3),
       )(rc, 0, depth_data, 5.0)
 
-    self.assertEqual(depth.shape, (nworld, height, width))
+    self.assertEqual(depth.shape, (nworld, height, width, 1))
 
 
 if __name__ == '__main__':