chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix extra dim for depth, matches rgb output shape.

Browse Source 
PiperOrigin-RevId: 878098932
Change-Id: Ifcc4ba39dbee1beeea88fda44c3d20b88179705e
This commit is contained in:
Baruch Tabanpour
2026-03-03 13:26:00 -08:00
committed by Copybara-Service
parent 490c1f4126
commit a8d6b0c848
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/render_util.py
+1 -1
View File
@@ -110,4 +110,4 @@ def get_depth(
) )
depth = jnp.clip(raw / depth_scale, 0.0, 1.0) 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
+2 -2
View File
@@ -94,7 +94,7 @@ class RenderUtilTest(absltest.TestCase):
rc, 0, depth_data, 5.0 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): def test_get_depth_vmap(self):
nworld, width, height = 3, 4, 4 nworld, width, height = 3, 4, 4
@@ -111,7 +111,7 @@ class RenderUtilTest(absltest.TestCase):
static_argnums=(0, 1, 3), static_argnums=(0, 1, 3),
)(rc, 0, depth_data, 5.0) )(rc, 0, depth_data, 5.0)
self.assertEqual(depth.shape, (nworld, height, width)) self.assertEqual(depth.shape, (nworld, height, width, 1))
if __name__ == '__main__': if __name__ == '__main__':