Expose MJX light_active model field

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,49 @@ 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_replace_light_active(self):
+    """Tests that light_active can be replaced like other public fields."""
+    m = mujoco.MjModel.from_xml_string(_LIGHTS)
+    mx = mjx.put_model(m, impl='jax')
+
+    light_active = jp.array([True, False])
+    mx_replaced = mx.replace(light_active=light_active)
+
+    np.testing.assert_array_equal(
+        np.asarray(mx_replaced.light_active), np.asarray(light_active)
+    )
+    np.testing.assert_array_equal(np.asarray(mx.light_active), m.light_active)
+
+  def test_vmap_light_active_in_axes(self):
+    """Tests that light_active participates in batched model pytrees."""
+    m = mujoco.MjModel.from_xml_string(_LIGHTS)
+    mx = mjx.put_model(m, impl='jax')
+    light_active = jp.array([[True, False], [False, True], [True, True]])
+    mx_batched = mx.replace(light_active=light_active)
+
+    in_axes = jax.tree_util.tree_map(lambda x: None, mx)
+    in_axes = in_axes.replace(light_active=0)
+    mapped = jax.vmap(lambda model: model.light_active, in_axes=(in_axes,))(
+        mx_batched
+    )
+
+    np.testing.assert_array_equal(np.asarray(mapped), np.asarray(light_active))
+
   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

@@ -14,17 +14,19 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
+import warp as wp
+
 from mujoco.mjx._src import types
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 from mujoco.mjx.warp import ffi
 from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.render_context import RenderContextPytree
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
-import warp as wp
-
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,6 +50,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
@@ -149,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',
@@ -172,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

@@ -15,14 +15,18 @@
 """MJX Warp types.
 DO NOT EDIT. This file is auto-generated.
 """
+
 import dataclasses
 import typing
 from typing import Tuple
+
 import jax
 from jax import tree_util
 from jax.interpreters import batching
-from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
+
+from mujoco.mjx._src import dataclasses as mjx_dataclasses
+
 if typing.TYPE_CHECKING:
   GraphMode = int
 
@@ -33,6 +37,7 @@ if typing.TYPE_CHECKING:
 else:
   try:
     from warp._src.jax_experimental.ffi import GraphMode
+
     from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 
     Callback = mjwp_types.Callback
@@ -41,6 +46,7 @@ else:
     Callback = None
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
+
 @dataclasses.dataclass(frozen=True)
 @tree_util.register_pytree_node_class
 class TileSet:
@@ -52,6 +58,7 @@ class TileSet:
     adr: address of each tile in the set
     size: size of all the tiles in this set
   """
+
   adr: np.ndarray
   size: int
 
@@ -72,6 +79,7 @@ class BlockDim:
 
   TODO(team): experimental and may be removed
   """
+
   actuator_velocity: int
   cholesky_factorize: int
   cholesky_factorize_solve: int
@@ -102,10 +110,13 @@ class BlockDim:
 
 class StatisticWarp(PyTreeNode):
   """Derived fields from Statistic."""
+
   meaninertia: jax.Array
 
+
 class OptionWarp(PyTreeNode):
   """Derived fields from Option."""
+
   broadphase: int
   broadphase_filter: int
   ccd_iterations: int
@@ -120,8 +131,10 @@ class OptionWarp(PyTreeNode):
   sdf_initpoints: int
   sdf_iterations: int
 
+
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
+
   M_colind: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
@@ -188,7 +201,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
   mapM2M: np.ndarray
@@ -283,8 +295,10 @@ class ModelWarp(PyTreeNode):
   wrap_site_adr: np.ndarray
   wrap_site_pair_adr: np.ndarray
 
+
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
+
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -363,6 +377,8 @@ class DataWarp(PyTreeNode):
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
+
+
 DATA_NON_VMAP = {
     'contact__dim',
     'contact__dist',
@@ -390,6 +406,7 @@ DATA_NON_VMAP = {
     'nworld',
 }
 
+
 def _to_elt(cont, _, d, axis):
   return DataWarp(**{
       f.name: (