From a8d6b0c8487fff528d679132be0b49560619920e Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 3 Mar 2026 13:26:00 -0800 Subject: [PATCH] Fix extra dim for depth, matches rgb output shape. PiperOrigin-RevId: 878098932 Change-Id: Ifcc4ba39dbee1beeea88fda44c3d20b88179705e --- mjx/mujoco/mjx/_src/render_util.py | 2 +- mjx/mujoco/mjx/_src/render_util_test.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mjx/mujoco/mjx/_src/render_util.py b/mjx/mujoco/mjx/_src/render_util.py index 0ef94281..d0bd1110 100644 --- a/mjx/mujoco/mjx/_src/render_util.py +++ b/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) diff --git a/mjx/mujoco/mjx/_src/render_util_test.py b/mjx/mujoco/mjx/_src/render_util_test.py index 46c74420..7549b0af 100644 --- a/mjx/mujoco/mjx/_src/render_util_test.py +++ b/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__':