diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 21e1910c..e086c38c 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -113,6 +113,15 @@ _SIMPLE_BODY = """ """ +_LIGHTS = """ + + + + + + +""" + 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( diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index bc51e6e5..1f3cf7b4 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/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 diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py index cf6754a1..92311b58 100644 --- a/mjx/mujoco/mjx/warp/render.py +++ b/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, diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py index f64b6fba..ce786ddc 100644 --- a/mjx/mujoco/mjx/warp/types.py +++ b/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: (