Merge pull request #3278 from tkelestemur:tarik/mjx-light-active-public-api

PiperOrigin-RevId: 924570592
Change-Id: I86d506e9f03d3230eaaca7061f6794ecdbfe164f

mjx/mujoco/mjx/_src/io_test.py

@@ -113,6 +113,15 @@ _SIMPLE_BODY = """
   </mujoco>
 """
 
+_LIGHTS = """
+  <mujoco>
+    <worldbody>
+      <light name="inactive" active="false"/>
+      <light name="active" active="true"/>
+    </worldbody>
+  </mujoco>
+"""
+
 
 def _get_name_from_path(path: jax.tree_util.KeyPath) -> str:
   """Returns a flattened name from a jax.tree_util.KeyPath."""
@@ -203,6 +212,21 @@ class ModelIOTest(parameterized.TestCase):
     np.testing.assert_equal(mx.wrap_objid, m.wrap_objid)
     np.testing.assert_almost_equal(mx.wrap_prm, m.wrap_prm)
 
+  @parameterized.parameters('jax', 'warp', 'cpp')
+  def test_put_model_light_active(self, impl):
+    """Tests that light_active is copied to the public Model field."""
+    if impl == 'warp' and not mjxw.WARP_INSTALLED:
+      self.skipTest('Warp not installed.')
+
+    m = mujoco.MjModel.from_xml_string(_LIGHTS)
+    mx = mjx.put_model(m, impl=impl)
+
+    self.assertIn('light_active', [f.name for f in mx.fields()])
+    np.testing.assert_array_equal(
+        np.asarray(mx.light_active),
+        m.light_active.astype(np.asarray(mx.light_active).dtype),
+    )
+
   def test_fluid_params(self):
     """Test that has_fluid_params is set when fluid params are present."""
     m = mjx.put_model(

mjx/mujoco/mjx/_src/types.py

@@ -747,6 +747,7 @@ class Model(PyTreeNode):
   cam_intrinsic: jax.Array
   light_mode: np.ndarray
   light_type: jax.Array
+  light_active: jax.Array
   light_castshadow: jax.Array
   light_pos: jax.Array
   light_dir: jax.Array

mjx/mujoco/mjx/warp/render.py

@@ -152,6 +152,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
           'geom_size',
           'geom_xmat',
           'geom_xpos',
+          'light_active',
           'light_castshadow',
           'light_type',
           'mat_rgba',
@@ -175,7 +176,7 @@ def _render_jax_impl(m: types.Model, d: types.Data, ctx: RenderContextPytree):
       m.geom_rgba,
       m.geom_size,
       m.geom_type,
-      m._impl.light_active,
+      m.light_active,
       m.light_castshadow,
       m.light_type,
       m.mat_rgba,

mjx/mujoco/mjx/warp/types.py

@@ -240,7 +240,6 @@ class ModelWarp(PyTreeNode):
   is_sparse: bool
   jnt_limited_ball_adr: np.ndarray
   jnt_limited_slide_hinge_adr: np.ndarray
-  light_active: jax.Array
   light_bodyid: np.ndarray
   light_targetbodyid: np.ndarray
   mapD2M: np.ndarray