diff --git a/mjx/mujoco/mjx/codegen/README.md b/mjx/mujoco/mjx/codegen/README.md new file mode 100644 index 00000000..3161b534 --- /dev/null +++ b/mjx/mujoco/mjx/codegen/README.md @@ -0,0 +1,23 @@ +# MJX Warp Codegen + +Generates the MJX-Warp shim layer in `mujoco/mjx/warp/` by reading the vendored +`mujoco_warp` source in `mujoco/mjx/third_party/mujoco_warp/`. + +## Setup + +From the root `mjx/` directory, once you install [`uv`](https://docs.astral.sh/uv/getting-started/installation/), install the latest MuJoCo and local MJX: + +```bash +uv venv .venv --default-index https://pypi.org/simple +source .venv/bin/activate +uv pip install --upgrade --force-reinstall mujoco --default-index https://pypi.org/simple --extra-index-url https://py.mujoco.org/ +uv pip install -e ".[warp,dev]" --default-index https://pypi.org/simple +``` + +## Run codegen + +From the root `mjx/` directory: + +```bash +bash mujoco/mjx/codegen/update_for_mujoco_warp.sh +``` diff --git a/mjx/mujoco/mjx/codegen/__init__.py b/mjx/mujoco/mjx/codegen/__init__.py new file mode 100644 index 00000000..100b5347 --- /dev/null +++ b/mjx/mujoco/mjx/codegen/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/mjx/mujoco/mjx/codegen/file.py b/mjx/mujoco/mjx/codegen/file.py new file mode 100644 index 00000000..06d6e6d1 --- /dev/null +++ b/mjx/mujoco/mjx/codegen/file.py @@ -0,0 +1,98 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""File tools.""" + +import ast +import os +import subprocess +from typing import Dict +from absl import logging +from etils import epath + +LICENSE_TEXT = """ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" + + +def get_base_path() -> epath.Path: + """Resolves the base workspace path.""" + base_path = os.environ.get('BUILD_WORKSPACE_DIRECTORY') + if base_path: + return epath.Path(base_path) + # Assume this file is at /mujoco/mjx/codegen/file.py. + root = epath.Path(os.path.abspath(__file__)).parents[3] + if not (root / 'mujoco' / 'mjx' / 'codegen').is_dir(): + raise RuntimeError( + f'Unexpected codegen layout, resolved root: {root}' + ) + return root + + +def format_file(target_fpath: epath.Path): + """Formats a Python file.""" + logging.info('Running pyink on: %s', target_fpath) + subprocess.run( + ['pyink', str(target_fpath)], + check=True, + text=True, + capture_output=True, + ) + logging.info('Running isort on: %s', target_fpath) + subprocess.run( + ['isort', str(target_fpath)], + check=True, + text=True, + capture_output=True, + ) + + +def write_license(target_fpath: epath.Path): + """Writes license to the target file.""" + src = target_fpath.read_text() + target_fpath.write_text(LICENSE_TEXT + src) + + +def get_cls_type_annotations(src: str) -> Dict[str, Dict[str, str]]: + """Return classes with their field annotation strings from source code.""" + ret = {} + tree = ast.parse(src) + + class Visitor(ast.NodeVisitor): + + def visit_ClassDef(self, node: ast.ClassDef): # pylint: disable=invalid-name + class_name = node.name + ret[class_name] = {} + for item in node.body: + if not isinstance(item, ast.AnnAssign): + continue + field_name = item.target.id # pytype: disable=attribute-error + annotation_str = ast.unparse(item.annotation).strip() + ret[class_name][field_name] = annotation_str + + Visitor().visit(tree) + return ret diff --git a/mjx/mujoco/mjx/codegen/generate_warp_shim.py b/mjx/mujoco/mjx/codegen/generate_warp_shim.py new file mode 100644 index 00000000..84c63925 --- /dev/null +++ b/mjx/mujoco/mjx/codegen/generate_warp_shim.py @@ -0,0 +1,536 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Creates a shim between JAX and Warp for a given function.""" + +import enum +import inspect +import re +from typing import Dict, List, Sequence + +from absl import app +from absl import flags +from absl import logging +from etils import epath +from mujoco.mjx._src import types as mjx_types +from mujoco.mjx.codegen import file +from mujoco.mjx.codegen import trace +import jax +from mujoco.mjx.third_party import mujoco_warp # pylint: disable=unused-import + + +_MJWARP_FUNCTION = flags.DEFINE_string( + 'mjwarp_function', + 'third_party/py/mujoco_warp/_src/smooth.py:kinematics', + 'Function to create the shim for.', +) +_MJWARP_TYPES = flags.DEFINE_string( + 'mjwarp_types', + 'third_party/py/mujoco_warp/_src/types.py', + 'Path to the mjwarp types file.', +) +_MJX_WARP_OUTPUT_PATH = flags.DEFINE_string( + 'mjx_warp_output_path', + 'third_party/py/mujoco/mjx/warp', + 'Path to the output file.', +) +_ONLY_PUBLIC_OUTPUT_FIELDS = flags.DEFINE_bool( + 'only_public_output_fields', + False, + 'Whether to keep only public fields in the output.', +) +_APPEND_TO_OUTPUT_FILE = flags.DEFINE_bool( + 'append_to_output_file', + False, + 'Whether to append to the output file.', +) + +_RENDER_CONTEXT_BUFFER_NAME = '_MJX_RENDER_CONTEXT_BUFFERS' + + +def _clean_type(type_: str): + # check for enums + if type_ in [ + name + for name, obj in inspect.getmembers(mujoco_warp, inspect.isclass) + if issubclass(obj, (enum.IntEnum, enum.IntFlag)) + ]: + return 'int' + + types_to_prefix = ( + 'vec5', + 'vec8', + 'vec8i', + 'vec10', + 'vec10f', + 'vec11', + 'TileSet', + 'BlockDim', + 'vec_pluginattr', + ) + m = re.match(r'array\((.*)\)', type_) + if m: # match custom mujoco_warp array annotation types + args_str = m.group(1) + args = [a.strip() for a in args_str.split(',')] + ndim, dtype = len(args) - 1, args[-1] + + dims = {1: '', 2: '2d', 3: '3d', 4: '4d'} + if ndim not in dims: + raise ValueError(f'Unsupported array ndim: {ndim} for type: {dtype}') + + type_ = f'wp.array{dims[ndim]}[{dtype}]' + + for t in types_to_prefix: + type_ = re.sub(rf'\b{t}\b', f'mjwp_types.{t}', type_) + return type_ + + +def _get_stage_fields( + field_usage: trace.FieldUsage, +) -> tuple[list[str], list[str]]: + """Returns stage_in and stage_out fields after tracing. + + stage_in: + * Model/ModelWarp jax.Array input fields + * Data jax.Array input fields + * Option/OptionWarp jax.Array input fields + + stage_out: + * Data jax.Array output fields + + Args: + field_usage: FieldUsage object + + Returns: + A tuple of (stage_in, stage_out) field name lists. + """ + stage_in = [] + stage_out = [] + + def is_jax_array(cls, field): + if cls is None: + return False + return cls.__annotations__.get(field) is jax.Array + + ModelWarp = getattr(mjx_types, 'ModelWarp', None) + OptionWarp = getattr(mjx_types, 'OptionWarp', None) + + # stage_in: Model/ModelWarp jax.Array input fields + for field in field_usage.model_fields: + if is_jax_array(mjx_types.Model, field): + stage_in.append(field) + elif is_jax_array(ModelWarp, field): + stage_in.append(field) + # stage_in: Option/OptionWarp jax.Array input fields + elif field.startswith('opt__'): + sub_field = field.split('opt__')[-1] + if is_jax_array(mjx_types.Option, sub_field): + stage_in.append(field) + elif is_jax_array(OptionWarp, sub_field): + stage_in.append(field) + + # stage_in: Data jax.Array input fields + for field in field_usage.data_fields: + if is_jax_array(mjx_types.Data, field): + stage_in.append(field) + + # stage_out: Data jax.Array output fields + for field in field_usage.data_out_fields: + if is_jax_array(mjx_types.Data, field): + stage_out.append(field) + + return sorted(stage_in), sorted(stage_out) + + +def _top_level_imports(field_usage: trace.FieldUsage): + """Returns top-level imports.""" + imports = ''' +"""DO NOT EDIT. This file is auto-generated.""" +import dataclasses +import functools +from mujoco.mjx._src import types +from mujoco.mjx.warp import ffi +import mujoco.mjx.third_party.mujoco_warp as mjwarp +import warp as wp +import jax +from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types +''' + + if field_usage.render_context_in_caller: + imports += ( + """ +from mujoco.mjx.warp.render_context import """ + + _RENDER_CONTEXT_BUFFER_NAME + + """ +from mujoco.mjx.warp.render_context import RenderContextPytree +""" + ) + + return imports + + +def _global_assignments(): + """Returns global assignments.""" + assignments = '' + for attr, cls in ( + ('_m', 'mjwarp.Model'), + ('_d', 'mjwarp.Data'), + ('_o', 'mjwarp.Option'), + ('_s', 'mjwarp.Statistic'), + ('_c', 'mjwarp.Contact'), + ('_e', 'mjwarp.Constraint'), + ('_cb', 'mjwp_types.Callback'), + ): + assignments += ( + f'{attr} = {cls}(**{{f.name: None for f in dataclasses.fields({cls}) if' + ' f.init})\n' + ) + return assignments + + +def _warp_function( + fn_name: str, + field_usage: trace.FieldUsage, + mjwarp_field_info: Dict[str, trace.FieldInfo], + mjx_warp_field_info: Dict[str, trace.FieldInfo], +): + """Returns warp function arguments, assignments, and call.""" + # create warp function. + fn_args_model, fn_assignments = [('nworld: int,', (-1, ''))], [] + if field_usage.model_fields: + for f in field_usage.model_fields: + if f not in mjwarp_field_info: + raise AssertionError(f'Field {f} not found in mjwarp_field_info.') + info = mjwarp_field_info[f] + expected_type = _clean_type(info.expected_type) + fn_args_model.append((f'{f}: {expected_type},', info.param_order)) + fn_assignments.append(f' _m.{f.replace('__', '.')} = {f}') + fn_args_model = sorted(fn_args_model, key=lambda x: x[1]) + fn_args_model = ['# Model'] + [f[0] for f in fn_args_model] + + fn_args_data = [] + if field_usage.data_fields: + for f in field_usage.data_fields: + if f not in mjwarp_field_info: + raise AssertionError(f'Field {f} not found in mjwarp_field_info.') + if f == 'nworld': + continue # this gets set manually below + j = trace.FieldInfo(f, 'jax.Array', (0, '')) + is_jax_arr = mjx_warp_field_info.get(f, j).expected_type == 'jax.Array' + is_out = is_jax_arr + info = mjwarp_field_info[f] + param_order = info.param_order + expected_type = _clean_type(info.expected_type) + fn_args_data.append((f'{f}: {expected_type},', (is_out, param_order))) + fn_assignments.append(f' _d.{f.replace('__', '.')} = {f}') + fn_args_data = sorted(fn_args_data, key=lambda x: x[1]) + fn_args_data = ['# Data'] + [f[0] for f in fn_args_data] + + fn_assignments.append(' _d.nworld = nworld') + + render_context_args = [] + render_context_call_arg = '' + if field_usage.render_context_in_caller: + render_context_args = ['# Registry', 'rc_id: int,'] + render_context_call_arg = ', render_context' + fn_assignments.append( + f' render_context = {_RENDER_CONTEXT_BUFFER_NAME}[(rc_id, wp.get_device().ordinal)]' + ) + + if fn_name == 'render': + render_context_args.append('rgb: wp.array2d[wp.uint32],') + render_context_args.append('depth: wp.array2d[wp.float32],') + fn_assignments.append(' render_context.rgb_data = rgb') + fn_assignments.append(' render_context.depth_data = depth') + else: + fn_assignments.append(' dummy.zero_()') + + fn_call = f'mjwarp.{fn_name}(_m, _d{render_context_call_arg})' + fn_args_raw = fn_args_model + fn_args_data + render_context_args + + # create a dummy output if there are no output fields + needs_dummy_output = not field_usage.data_out_fields + if needs_dummy_output and fn_name != 'render': + fn_args_raw.append('# Dummy output') + fn_args_raw.append('dummy: wp.array[int],') + + return fn_args_raw, fn_assignments, fn_call + + +def _jax_shim_fn( + fn_name: str, + field_usage: trace.FieldUsage, + warp_fn_args: List[str], + mjwarp_field_info: Dict[str, trace.FieldInfo], +): + """Generates a JAX shim for the Warp function.""" + num_outputs = 0 + output_dims = [] + jax_args = [] + tree_replace = [] + in_out_argnames = [] + has_side_effect = False + stage_in_fields, stage_out_fields = _get_stage_fields(field_usage) + stage_in_argnames = [f"'{f}'" for f in stage_in_fields] + stage_out_argnames = [f"'{f}'" for f in stage_out_fields] + + for arg in warp_fn_args: + if 'nworld' in arg: + jax_args.append('d.qpos.shape[0]') + continue + + if arg in ('rc_id', 'dummy'): + continue + + if arg in ('rgb', 'depth') and fn_name == 'render': + num_outputs += 1 + continue + + arg_jax = arg + if mjwarp_field_info[arg].param_source == 'Data': + if arg.split('__')[0] not in mjx_types.Data.__annotations__: + arg_jax = f'_impl.{arg}' + arg_jax = 'd.' + arg_jax + elif mjwarp_field_info[arg].param_source == 'Model': + arg_jax = ( + arg.replace('__', '.') + if arg.startswith('opt') or arg.startswith('stat') + else arg + ) + public_field = arg.split('__')[0] in mjx_types.Model.__annotations__ + if not public_field: + arg_jax = f'_impl.{arg_jax}' + if ( + arg.startswith('opt') + and arg.split('__')[-1] not in mjx_types.Option.__annotations__ + ): + arg_jax = arg_jax.replace('opt', 'opt._impl') + arg_jax = 'm.' + arg_jax + else: + raise ValueError( + f'Unknown param source: {mjwarp_field_info[arg].param_source}' + ) + + if arg in field_usage.data_out_fields: + # all out fields are in_out, since JAX already allocated them + in_out_argnames.append(f"'{arg}'") + num_outputs += 1 + output_dims.append(f"'{arg}': {arg_jax}.shape") + + if '_impl' not in arg_jax or not _ONLY_PUBLIC_OUTPUT_FIELDS.value: + tree_replace.append(f'"{arg_jax[2:]}": out[{num_outputs - 1}]') + + if arg == 'geom_dataid': + jax_args.append(f'jax.numpy.expand_dims({arg_jax}, 0)') + else: + jax_args.append(arg_jax) + + if field_usage.render_context_in_caller: + jax_args.append('ctx.key') + + needs_dummy_output = not field_usage.data_out_fields + if needs_dummy_output and fn_name != 'render': + num_outputs = 1 + output_dims = ["'dummy': (d.qpos.shape[0],)"] + has_side_effect = True + + if fn_name == 'render': + output_dims = [ + "'rgb': render_ctx.rgb_data_shape", + "'depth': render_ctx.depth_data_shape", + ] + tree_replace = [] + + render_ctx_param = ( + 'ctx: RenderContextPytree' if field_usage.render_context_in_caller else '' + ) + fn_args = ['m: types.Model', 'd: types.Data'] + + if render_ctx_param: + fn_args.append(render_ctx_param) + + return ( + fn_args, + jax_args, + output_dims, + num_outputs, + tree_replace, + in_out_argnames, + stage_in_argnames, + stage_out_argnames, + has_side_effect, + ) + + +def create_jax_warp_shim( + fn_name: str, + field_usage: trace.FieldUsage, + mjwarp_field_info: Dict[str, trace.FieldInfo], + mjx_warp_field_info: Dict[str, trace.FieldInfo], + out_fpath: epath.Path, +): + """Creates a JAX-wrapped MJWarp function.""" + src = '' + old_src = ( + out_fpath.read_text() + if out_fpath.exists() + else '' + if out_fpath.exists() + else '' + ) + + # create top-level imports. + if not _APPEND_TO_OUTPUT_FILE.value: + src += _top_level_imports(field_usage) + '\n\n' + + # create global assignments. + assignments = _global_assignments() + already_in_src = re.sub(r'\s+', '', assignments) in re.sub( + r'\s+', '', old_src + ) + if not already_in_src or not _APPEND_TO_OUTPUT_FILE.value: + src += assignments + + # create warp function. + fn_args_raw, fn_assignments, fn_call = _warp_function( + fn_name, field_usage, mjwarp_field_info, mjx_warp_field_info + ) + fn_args_raw_str = '\n'.join([' ' + arg for arg in fn_args_raw]) + warp_fn_args = [arg.split(':')[0] for arg in fn_args_raw if '#' not in arg] # pytype: disable=attribute-error + + src += f""" +@ffi.format_args_for_warp +def _{fn_name}_shim( +{fn_args_raw_str} +): + _m.stat = _s + _m.opt = _o + _m.callback = _cb + _d.efc = _e + _d.contact = _c +{'\n'.join(fn_assignments)} + {fn_call} + """ + src += '\n\n' + + # create private jax function. + ( + fn_args, + jax_args, + output_dims, + num_outputs, + tree_replace, + in_out_argnames, + stage_in_argnames, + stage_out_argnames, + has_side_effect, + ) = _jax_shim_fn(fn_name, field_usage, warp_fn_args, mjwarp_field_info) + render_ctx_line = '' + return_stmt = 'return d' + if fn_name == 'render': + render_ctx_line = f' render_ctx = _MJX_RENDER_CONTEXT_BUFFERS[(ctx.key, None)]\n' + return_stmt = 'return out' + output_dims_str = '{' + ','.join(output_dims) + '}' + data_tree_replace = f"d = d.tree_replace({{ {','.join(tree_replace)} }})" + src += f""" +def _{fn_name}_jax_impl({','.join(fn_args)}): +{render_ctx_line} output_dims = {output_dims_str} + jf = ffi.jax_callable_variadic_tuple( + _{fn_name}_shim, num_outputs={num_outputs}, + output_dims=output_dims, + vmap_method=None, + in_out_argnames=set([{','.join(in_out_argnames)}]), + stage_in_argnames=set([{','.join(stage_in_argnames)}]), + stage_out_argnames=set([{','.join(stage_out_argnames)}]), + graph_mode=m.opt._impl.graph_mode, + has_side_effect={has_side_effect}, + ) + out = jf({','.join(jax_args)}) + {data_tree_replace} + {return_stmt} +""" + src += '\n' + + # create public jax functions. + fn_args_no_annotation = [arg.split(':')[0] for arg in fn_args] + fn_call_str = ','.join(fn_args_no_annotation) + + marshal_decorator = '@ffi.marshal_jax_warp_callable' + marshal_vmap_decorator = '@ffi.marshal_custom_vmap' + vmap_return_stmt = f'd = {fn_name}({fn_call_str})\n return d, is_batched[1]' + if fn_name == 'render': + marshal_decorator = ( + '@functools.partial(' + 'ffi.marshal_jax_warp_callable, tree_map_output=True)' + ) + marshal_vmap_decorator = ( + '@functools.partial(ffi.marshal_custom_vmap, tree_map_output=True)' + ) + vmap_return_stmt = ( + f'out = {fn_name}({fn_call_str})\n return out, [True, True]' + ) + + src += f""" +@jax.custom_batching.custom_vmap +{marshal_decorator} +def {fn_name}({','.join(fn_args)}): + return _{fn_name}_jax_impl({','.join(fn_args_no_annotation)}) +@{fn_name}.def_vmap +{marshal_vmap_decorator} +def {fn_name}_vmap(unused_axis_size, is_batched, {','.join(fn_args)}): + {vmap_return_stmt} +""" + src += '\n' + + if _APPEND_TO_OUTPUT_FILE.value: + src = old_src + '\n\n' + src + + out_fpath.write_text(src) + + +def main(argv: Sequence[str]) -> None: + del argv + logging.set_verbosity(logging.DEBUG) + + # Get mjwarp field annotations. + fpath = epath.Path(_MJWARP_TYPES.value) + mjwarp_field_info = trace.get_mjwarp_field_info( + fpath.read_text(), file.get_cls_type_annotations + ) + + # Trace function to get field usage. + fpath, fn_name = _MJWARP_FUNCTION.value.split(':') + field_usage = trace.trace_function(fpath, fn_name, mjwarp_field_info) + + base_path = file.get_base_path() + types_fpath = base_path / _MJX_WARP_OUTPUT_PATH.value / 'types.py' + mjx_warp_field_info = trace.get_mjx_warp_field_info( + types_fpath.read_text(), file.get_cls_type_annotations + ) + + base_path = file.get_base_path() + target_fpath = ( + base_path / _MJX_WARP_OUTPUT_PATH.value / epath.Path(fpath).name + ) + create_jax_warp_shim( + fn_name, field_usage, mjwarp_field_info, mjx_warp_field_info, target_fpath + ) + + if not _APPEND_TO_OUTPUT_FILE.value: + file.write_license(target_fpath) + file.format_file(target_fpath) + + +if __name__ == '__main__': + app.run(main) diff --git a/mjx/mujoco/mjx/codegen/generate_warp_types.py b/mjx/mujoco/mjx/codegen/generate_warp_types.py new file mode 100644 index 00000000..374e7464 --- /dev/null +++ b/mjx/mujoco/mjx/codegen/generate_warp_types.py @@ -0,0 +1,570 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Generate types for MJX warp integration.""" + +import ast +import dataclasses +import enum +import logging +import typing +from typing import Any, Callable, Dict, List, Optional, Set + +from absl import app +from absl import flags +from etils import epath +import mujoco +from mujoco.mjx.codegen import file +import mujoco.mjx.third_party.mujoco_warp as mjwarp +import numpy as np +import warp as wp +from mujoco.mjx.third_party.warp._src.jax_experimental import ffi + + +_MJX_WARP_TYPES_OUT_FPATH = flags.DEFINE_string( + 'mjx_warp_types_out_path', + 'third_party/py/mujoco/mjx/warp/types.py', + 'Path to write the mjWarp types into.', +) + +_MJX_TYPES_PATH = flags.DEFINE_string( + 'mjx_types_path', + 'third_party/py/mujoco/mjx/_src/types.py', + 'Path to read the MJX types from.', +) + +_DATA_SHAPE_PROPERTY_FIELD = 'cacc' +_DUMMY_XML = """ + + + + + + + + + + + + +""" + + +def _to_py_string(value, indent=0): + """Converts a dictionary/set/tuple/type to a Python code string.""" + indent_str = ' ' * indent + next_indent_str = ' ' * (indent + 1) + if isinstance(value, tuple): + items = [_to_py_string(item, indent) for item in value] + return f'({', '.join(items)})' + + if isinstance(value, type): + if value.__module__ == 'builtins': + return value.__name__ + return f'{value.__module__}.{value.__name__}' + + if isinstance(value, dict): + items = [ + f'\n{next_indent_str}{repr(k)}: {_to_py_string(v, indent + 1)}' + for k, v in sorted(value.items(), key=lambda x: x[0]) + ] + return f'{{{','.join(items)}\n{indent_str}}}' + + if isinstance(value, set): + items = sorted([_to_py_string(item, indent) for item in value]) + items = [f'\n{next_indent_str}{item}' for item in items] + return f'{{{",".join(items)}\n{indent_str}}}' + + return repr(value) + + +def _ast_parse_type(type_repr: str) -> ast.expr: + """Parses a string representation of a type into an AST node.""" + try: + return ast.parse(type_repr, mode='eval').body + except SyntaxError as e: + raise ValueError(f'Failed to parse type repr "{type_repr}": {e}') from e + + +def _get_target_annotation_node( + key: str, + target_annotations: Dict[str, Any], +) -> ast.expr: + """Determines the AST node for the target type annotation for MJX.""" + annotation = target_annotations.get(key) + if annotation == np.ndarray: + return _ast_parse_type('np.ndarray') + + if (isinstance(annotation, wp.array) or + type(annotation).__name__ == '_ArrayAnnotation'): + return _ast_parse_type('jax.Array') + + if annotation in (int, float, bool): + return _ast_parse_type(annotation.__name__) + + if annotation is ffi.GraphMode: + return _ast_parse_type('GraphMode') + + if isinstance(annotation, type) and issubclass(annotation, enum.Enum): + return _ast_parse_type('int') + + if dataclasses.is_dataclass(annotation): + return _ast_parse_type(annotation.__name__) + + is_tuple = typing.get_origin(annotation) == tuple + if is_tuple and typing.get_args(annotation)[1] != ...: + raise NotImplementedError( + 'Only variadic tuples are supported. Got annotation type' + f' {annotation} for key {key}.' + ) + + if is_tuple and typing.get_args(annotation)[0] in (int, float, bool): + type_ = typing.get_args(annotation)[0].__name__ + return _ast_parse_type(f'Tuple[{type_}, ...]') + + if is_tuple and ( + isinstance(typing.get_args(annotation)[0], wp.array) + or type(typing.get_args(annotation)[0]).__name__ == '_ArrayAnnotation' + ): + return _ast_parse_type('Tuple[np.ndarray, ...]') + + if is_tuple and dataclasses.is_dataclass(typing.get_args(annotation)[0]): + cls_type = typing.get_args(annotation)[0].__name__ + return _ast_parse_type(f'Tuple[{cls_type}, ...]') + + raise NotImplementedError( + f'Unhandled annotation type {annotation} for key {key}.' + ) + + +def _get_annotations_recursive( + annotations: Dict[str, Any], prefix: str = '' +) -> Dict[str, Any]: + """Recursively flattens type annotations, handling nested classes. + + Args: + annotations: A dictionary of type annotations (field_name: type). + prefix: The prefix to add to each key, for nested classes. + + Returns: + A dictionary of flattened annotations (e.g., 'contact__dist'). + """ + flattened = {} + for key, annotation in annotations.items(): + full_key = f'{prefix}{key}' + if hasattr( + annotation, '__annotations__' + ) and 'mujoco_warp' in annotation.__module__: + nested = _get_annotations_recursive( + dict(annotation.__annotations__), prefix=f'{full_key}__' + ) + flattened.update(nested) + else: + flattened[full_key] = annotation # Leaf node. + + return flattened + + +def _build_new_class_body_ast( + keys: Set[str], + cls_name: str, + target_annotations: Dict[str, Any], + shape_property: Optional[str] = None, + add_docstring: bool = True, +) -> List[ast.AST]: + """Builds the list of AST nodes for the new class body.""" + new_body_nodes: List[ast.AST] = [] + + if add_docstring: + docstring = f'Derived fields from {cls_name}.' + new_body_nodes.append(ast.Expr(value=ast.Constant(value=docstring))) + + # Sort keys alphabetically before creating AST nodes + sorted_keys = sorted(list(keys)) + for key in sorted_keys: + annotation_node = _get_target_annotation_node(key, target_annotations) + + new_body_nodes.append( + ast.AnnAssign( + target=ast.Name(id=key, ctx=ast.Store()), + annotation=annotation_node, + simple=1, # No value assignment + ) + ) + + if shape_property is not None: + property_string = ( + f'shape = property(lambda self: self.{shape_property}.shape)' + ) + shape_property_node = ast.parse(property_string).body[0] + new_body_nodes.append(shape_property_node) + + return new_body_nodes + + +def _write_class_in_file( + target_fpath: epath.Path, + target_cls_name: str, + target_base_name: str, + new_body_ast: List[ast.AST], +) -> None: + """Reads target file, writes the specified class, and saves.""" + target_code = target_fpath.read_text() + target_tree = ast.parse(target_code) + + if target_cls_name in target_code: + raise ValueError( + f'Class {target_cls_name} already exists in file: {target_fpath}' + ) + + new_class_def = ast.ClassDef( + name=target_cls_name, + bases=[ast.Name(id=target_base_name, ctx=ast.Load())], + body=new_body_ast, # pytype: disable=wrong-arg-types + decorator_list=[], + keywords=[], + type_params=[], + ) + target_tree.body.append(new_class_def) + + ast.fix_missing_locations(target_tree) + modified_code = ast.unparse(target_tree) + + logging.info('Writing modified code back to: %s', target_fpath) + with target_fpath.open('w') as f: + f.write(modified_code) + + logging.info('File successfully rewritten.') + + +def write_header(target_fpath: epath.Path): + """Writes imports and pre-defined class definitions to types.py.""" + header = ''' +"""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 + +if typing.TYPE_CHECKING: + GraphMode = int # Type alias for pytype. + @dataclasses.dataclass + class Callback: + pass +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 + except ImportError: + GraphMode = int # Fallback when warp not installed. + Callback = None + +PyTreeNode = mjx_dataclasses.PyTreeNode +''' + target_fpath.write_text(header) + + +_FLATTEN_UNFLATTEN = """ + # flatten/unflatten all fields for custom jax_callable, but prevent the parent + # PyTreeNode from putting these fields on device. + def tree_flatten(self): + children = list(getattr(self, k) for k in self.__dataclass_fields__) + return (children, None) + @classmethod + def tree_unflatten(cls, aux_data, children): + del aux_data + return cls(*children) +""" + + +def write_nested_dataclass(target_fpath: epath.Path, cls: Any): + new_class_body = _build_new_class_body_ast( + set(cls.__annotations__.keys()), + cls.__name__, + dict(cls.__annotations__), + add_docstring=False, + ) + cls_str = '\n'.join([' ' + ast.unparse(node) for node in new_class_body]) + cls_str = cls_str.replace('jax.Array', 'np.ndarray') + with target_fpath.open('a') as f: + f.write(f''' +@dataclasses.dataclass(frozen=True) +@tree_util.register_pytree_node_class +class {cls.__name__}: + """{cls.__doc__}""" +{cls_str} +{_FLATTEN_UNFLATTEN} +''') + + +def _get_meta_fields(cls_name: str) -> Set[str]: + """Returns the set of fields that should be meta-fields in the pytree.""" + m = mujoco.MjModel.from_xml_string(_DUMMY_XML) + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + + if cls_name == 'Data': + with wp.ScopedDevice('cpu'): + dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113) + cond_fn = lambda x: x.shape[0] not in {113, 1113} + dw_meta = _get_fields_with_cond(dw, cond_fn) + data_non_vmap = _get_non_vmap_data_fields() + return {k for k in dw_meta if k not in data_non_vmap} + + with wp.ScopedDevice('cpu'): + mw = mjwarp.put_model(m) + cond_fn = lambda x: not hasattr(x, '_is_batched') + mw_meta = _get_fields_with_cond(mw, cond_fn) + if cls_name == 'Model': + return mw_meta + + if cls_name == 'Option': + return {k[len('opt__') :] for k in mw_meta if k.startswith('opt')} + + if cls_name == 'Statistic': + return {k[len('stat__') :] for k in mw_meta if k.startswith('stat')} + + raise NotImplementedError(f'Unhandled class name {cls_name}.') + + +def write_core_cls( + cls_name: str, + target_fpath: epath.Path, + mjx_types_fpath: epath.Path, + flatten_fields: bool = False, + set_diff: bool = True, + extra_annotations: dict[str, type] | None = None, +): + """Writes a core API class (e.g. Model/Data/Option/Statistic).""" + cls = { + 'Model': mjwarp.Model, + 'Data': mjwarp.Data, + 'Option': mjwarp.Option, + 'Statistic': mjwarp.Statistic, + }[cls_name] + + annotations = dict(cls.__annotations__) # pytype: disable=attribute-error + if flatten_fields: + annotations = _get_annotations_recursive(annotations) + + meta_fields = _get_meta_fields(cls_name) + for k, v in annotations.items(): + if k not in meta_fields: + continue + if isinstance(v, wp.array) or type(v).__name__ == '_ArrayAnnotation': + annotations[k] = np.ndarray + + warp_keys = annotations.keys() + mjx_annotations = file.get_cls_type_annotations(mjx_types_fpath.read_text())[ + cls_name + ] + + keys = warp_keys + if set_diff: + # Take the set difference between warp and mjx annotation keys. + keys = warp_keys - mjx_annotations.keys() + + if extra_annotations: + annotations.update(extra_annotations) + keys = set(keys) | extra_annotations.keys() + + if not keys: + raise ValueError('No derived keys found') + + shape_property = None + if cls_name == 'Data': + shape_property = _DATA_SHAPE_PROPERTY_FIELD + + new_class_body = _build_new_class_body_ast( + keys, + cls_name, + annotations, + shape_property=shape_property, + ) + + _write_class_in_file( + target_fpath=target_fpath, + target_cls_name=f'{cls_name}Warp', + target_base_name='PyTreeNode', + new_body_ast=new_class_body, + ) + + +def _get_fields_with_cond( + d: Any, + cond_fn: Callable[[Any], bool], + s: Optional[Set[str]] = None, + prefix: str = '', + not_in: bool = False, + add_static: bool = False, +) -> Set[str]: + """Recursively finds fields given a condition on the leading dimensions.""" + if s is None: + s = set() + for f in dataclasses.fields(d): + attr = getattr(d, f.name) + if dataclasses.is_dataclass(f.type): + s = _get_fields_with_cond(attr, cond_fn, s, f.name + '__', not_in) + continue + if f.type in (int, float, bool) and add_static: + s.add(prefix + f.name) + continue + if not (isinstance(f.type, wp.array) or + type(f.type).__name__ == '_ArrayAnnotation'): + continue + if cond_fn(attr): + s.add(prefix + f.name) + return s + + +def _get_non_vmap_data_fields() -> Set[str]: + """Returns the fields that are not be vmapped but are still jax.Array.""" + m = mujoco.MjModel.from_xml_string(_DUMMY_XML) + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + + dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113) + cond_fn = lambda x: x.shape[0] not in {113} + non_vmap = _get_fields_with_cond(dw, cond_fn, add_static=True) + return non_vmap + + +def write_register_vmappable(target_fpath: epath.Path): + """Writes register_vmappable to types.py.""" + data_non_vmap = _get_non_vmap_data_fields() + with target_fpath.open('a') as f: + f.write('\nDATA_NON_VMAP =' + _to_py_string(data_non_vmap)) + f.write("""\n +def _to_elt(cont, _, d, axis): + return DataWarp(**{f.name: cont(getattr(d, f.name), axis) + if f.name not in DATA_NON_VMAP + else getattr(d, f.name) for f in DataWarp.fields()}) +def _from_elt(cont, axis_size, d, axis_dest): + return DataWarp(**{f.name: cont(axis_size, getattr(d, f.name), axis_dest) + if f.name not in DATA_NON_VMAP + else getattr(d, f.name) for f in DataWarp.fields()}) +batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None) +""") + + +def _is_ffi_compatible(wp_type: Any) -> bool: + """Returns True if the type is an array, scalar, or variadic tuple.""" + if (isinstance(wp_type, wp.array) or + type(wp_type).__name__ == '_ArrayAnnotation'): + return True + if wp_type in wp._src.types.value_types: + return True + if typing.get_origin(wp_type) is tuple: + return True + return False + + +def _to_jax_ndim(name: str, wp_type: Any) -> int: + if typing.get_origin(wp_type) is tuple: + if typing.get_args(wp_type)[1] != ...: + raise NotImplementedError('Only variadic tuples are supported.') + return -1 # signals that dim should be untouched in downstream code. + ffi_arg = ffi.FfiArg(name, wp_type) + return ffi_arg.jax_ndim + + +def write_ndim_annotations(target_fpath: epath.Path): + """Writes a dictionary of ndim for every warp field.""" + ndim_annotations = {} + for cls in [mjwarp.Model, mjwarp.Data, mjwarp.Option, mjwarp.Statistic]: + ndim_annotations[cls.__name__] = {} + annotations = _get_annotations_recursive(cls.__annotations__) + for name, type_ in annotations.items(): + if not _is_ffi_compatible(type_): + continue + ndim = _to_jax_ndim(name, type_) + ndim_annotations[cls.__name__][name] = ndim + + with target_fpath.open('a') as f: + f.write('\n_NDIM = ' + _to_py_string(ndim_annotations)) + + +def write_nworld_leading_dim(target_fpath: epath.Path): + """Writes a dictionary of which MJW fields are batched.""" + # TODO(btaba): check that batch fields have MJX jax.Array annotations, and + # that non-batch fields have np.ndarray annotations. Fail early. + + m = mujoco.MjModel.from_xml_string(_DUMMY_XML) + d = mujoco.MjData(m) + + batched = {} + for cls in [mjwarp.Model, mjwarp.Data, mjwarp.Option, mjwarp.Statistic]: + batched[cls.__name__] = {} + + if cls.__name__ == 'Data': + with wp.ScopedDevice('cpu'): + dw = mjwarp.put_data(m, d, nworld=113, naconmax=1113, njmax=1113) + cond_fn = lambda x: x.shape[0] == 113 + batched_fields = _get_fields_with_cond(dw, cond_fn) + else: + with wp.ScopedDevice('cpu'): + obj = mjwarp.put_model(m) + if cls.__name__ == 'Option': + obj = obj.opt + elif cls.__name__ == 'Statistic': + obj = obj.stat + cond_fn = lambda x: hasattr(x, '_is_batched') + batched_fields = _get_fields_with_cond(obj, cond_fn) + + all_annotations = _get_annotations_recursive(cls.__annotations__) + for name, type_ in all_annotations.items(): + if not _is_ffi_compatible(type_): + continue + batched[cls.__name__][name] = name in batched_fields + + with target_fpath.open('a') as f: + f.write('\n_BATCH_DIM = ' + _to_py_string(batched)) + + +def main(argv): + del argv + + base_path = file.get_base_path() + target_fpath = base_path / _MJX_WARP_TYPES_OUT_FPATH.value + mjx_types_fpath = base_path / _MJX_TYPES_PATH.value + + write_header(target_fpath) + # TODO(btaba): consider automated grabbing of nested dataclasses from mjwarp. + write_nested_dataclass(target_fpath, mjwarp._src.types.TileSet) + write_nested_dataclass(target_fpath, mjwarp._src.types.BlockDim) + + write_core_cls('Statistic', target_fpath, mjx_types_fpath, set_diff=False) + write_core_cls( + 'Option', target_fpath, mjx_types_fpath, + extra_annotations={'graph_mode': ffi.GraphMode}, + ) + write_core_cls('Model', target_fpath, mjx_types_fpath) + write_core_cls('Data', target_fpath, mjx_types_fpath, flatten_fields=True) + write_register_vmappable(target_fpath) + write_ndim_annotations(target_fpath) + write_nworld_leading_dim(target_fpath) + + file.write_license(target_fpath) + file.format_file(target_fpath) + + +if __name__ == '__main__': + app.run(main) diff --git a/mjx/mujoco/mjx/codegen/trace.py b/mjx/mujoco/mjx/codegen/trace.py new file mode 100644 index 00000000..5ba8359d --- /dev/null +++ b/mjx/mujoco/mjx/codegen/trace.py @@ -0,0 +1,318 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Static AST tracing to find MuJoCo Model and Data field usages.""" + +import ast +import dataclasses +import functools +import importlib +import importlib.util +import os +from typing import Dict, Optional, Sequence, Set, Tuple + +from absl import logging +from etils import epath +from mujoco.mjx.codegen import file + + +def _get_imported_module_names(fpath: epath.Path) -> Sequence[Tuple[str, str]]: + """Returns set of (fully qualified module_name, alias) tuples.""" + module_names = set() + tree = ast.parse(fpath.read_text(), filename=fpath) + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + module_names.add((alias.name, alias.name)) + elif isinstance(node, ast.ImportFrom): + if node.module: + for name in node.names: + name_ = name.name if not name.asname else name.asname + module_names.add((f'{node.module}.{name.name}', f'{name_}')) + return list(module_names) + + +def _resolve_module_name_to_fpath(fully_qualified_name: str) -> Optional[str]: + """Resolves a fully qualified name to its file path using importlib.""" + name_parts = fully_qualified_name.split('.') + for i in range(len(name_parts), 0, -1): + module_name_to_try = '.'.join(name_parts[:i]) + try: + spec = importlib.util.find_spec(module_name_to_try) + if spec and spec.origin and spec.origin != 'built-in': + return os.path.abspath(spec.origin) + except (ModuleNotFoundError, ImportError): + continue + + +def _get_imported_module_fpaths(fpath: epath.Path) -> Dict[str, str]: + """Returns the file paths of all imported modules.""" + all_imported_names = _get_imported_module_names(fpath) + + all_resolved_fpaths = {} + for fully_qualified_name, alias in all_imported_names: + fpath = _resolve_module_name_to_fpath(fully_qualified_name) + if fpath: + all_resolved_fpaths[alias] = fpath + return all_resolved_fpaths + + +@dataclasses.dataclass +class FieldInfo: + param_source: str + expected_type: str + param_order: tuple[int, str] + + +class _FunctionFieldUsageVisitor(ast.NodeVisitor): + """AST visitor to find attribute usages on 'm' and 'd' variables.""" + + def __init__( + self, + current_fpath: epath.Path, + visited_fns: Set[Tuple[str, str]], + mjwarp_field_info: Dict[str, FieldInfo], + ): + self.model_fields = set() + self.data_fields = set() + self.data_out_fields = set() + self._current_fpath = current_fpath.as_posix() + self._visited_fns = visited_fns + self._mjwarp_field_info = mjwarp_field_info + self._module_fpaths = _get_imported_module_fpaths(current_fpath) + self._in_outputs_context = False + + def visit_FunctionDef(self, node: ast.FunctionDef): + """Visits nested function definitions.""" + self.generic_visit(node) + + def add_field_usage(self, node: ast.Attribute, is_output: bool): + """Adds field to the appropriate set.""" + attr_parts = [] + curr_node = node + while isinstance(curr_node, ast.Attribute): + attr_parts.append(curr_node.attr) + curr_node = curr_node.value + + if isinstance(curr_node, ast.Name): + attr_parts.reverse() + full_attribute_str = '__'.join(attr_parts) + in_field_info = full_attribute_str in self._mjwarp_field_info + if in_field_info: + if curr_node.id == 'm': + self.model_fields.add(full_attribute_str) + if curr_node.id == 'd': + self.data_fields.add(full_attribute_str) + if curr_node.id == 'd' and is_output: + self.data_out_fields.add(full_attribute_str) + + def visit_Attribute(self, node: ast.Attribute): + self.add_field_usage(node, self._in_outputs_context) + self.generic_visit(node) + + def visit_keyword(self, node: ast.keyword): + """Visit a keyword argument node (e.g., outputs=[...]).""" + previous_in_outputs_context = self._in_outputs_context + if node.arg == 'outputs': + self._in_outputs_context = True + try: + self.generic_visit(node) + finally: + self._in_outputs_context = previous_in_outputs_context + + def recurse_trace(self, next_fpath: str, called_fn_name: str): + """Recursively trace into a function.""" + try: + field_usage = trace_function( + next_fpath, + called_fn_name, + self._mjwarp_field_info, + self._visited_fns, + ) + self.model_fields.update(field_usage.model_fields) + self.data_fields.update(field_usage.data_fields) + self.data_out_fields.update(field_usage.data_out_fields) + except ValueError as e: + logging.warning( + 'Could not trace function %s in %s: %s', + called_fn_name, + next_fpath, + e, + ) + + def visit_Call(self, node: ast.Call): + """Visit a function call node and recursively find all attribute usages.""" + if isinstance(node.func, ast.Name): + called_fn_name = node.func.id + key = (hash(self._current_fpath), called_fn_name) + if key not in self._visited_fns: + self._visited_fns.add(key) + next_fpath = self._module_fpaths.get( + called_fn_name, self._current_fpath + ) + self.recurse_trace(next_fpath, called_fn_name) + elif isinstance(node.func, ast.Attribute): + parts = [] + current = node.func + while isinstance(current, ast.Attribute): + parts.append(current.attr) + current = current.value + if isinstance(current, ast.Name): + parts.append(current.id) + parts = parts[::-1] + + if len(parts) == 2 and parts[0] == 'wp' and parts[1] == 'copy': + if len(node.args) != 2: + raise ValueError(f'wp.copy() must have 2 arguments, got {node.args}.') + out_node, in_node = node.args + self.add_field_usage(out_node, True) + self.add_field_usage(in_node, False) + return + + for arg in node.args: + if isinstance(arg, ast.Attribute): + self.add_field_usage(arg, is_output=True) + + called_fn_name = '.'.join(parts[1:]) + key = (hash(self._current_fpath), called_fn_name) + next_fpath = self._module_fpaths.get(parts[0]) + if next_fpath and key not in self._visited_fns: + self._visited_fns.add(key) + self.recurse_trace(next_fpath, called_fn_name) + + self.generic_visit(node) + + +@dataclasses.dataclass +class FieldUsage: + model_fields: Sequence[str] = dataclasses.field(default_factory=list) + data_fields: Sequence[str] = dataclasses.field(default_factory=list) + data_out_fields: Sequence[str] = dataclasses.field(default_factory=list) + render_context_in_caller: bool = False + + +def trace_function( + fpath: str, + fn: str, + mjwarp_field_info: Dict[str, FieldInfo], + visited_fns: Set[Tuple[str, str]] | None = None, +) -> FieldUsage: + """Traces the function statically to find usages of model and data fields.""" + base_path = file.get_base_path() + fpath = base_path / fpath + logging.info('Tracing function: "%s" in "%s"', fn, fpath) + + src = fpath.read_text() + parsed_ast = ast.parse(src, filename=str(fpath)) + + target_fn_nodes = ( + node + for node in parsed_ast.body + if isinstance(node, ast.FunctionDef) and node.name == fn + ) + target_fn_node = next(target_fn_nodes, None) + if not target_fn_node: + raise ValueError(f'Function "{fn}" not found in "{fpath}".') + + args_node = target_fn_node.args + args_tuple = tuple( + map(functools.partial(ast.get_source_segment, src), args_node.args) + ) + check_args_tuple = ( + ('m: Model', 'd: Data'), + ('m: types.Model', 'd: types.Data'), + ('m', 'd'), + ('m: Model',), + ('m: types.Model',), + ('m',), + ) + if ( + args_tuple[:2] not in check_args_tuple + and args_tuple[:1] not in check_args_tuple + ): + raise ValueError( + f'Function "{fn}" in "{fpath}" must have arguments in' + f' {check_args_tuple} got {args_tuple}.' + ) + + if visited_fns is None: + visited_fns = set() + + visitor = _FunctionFieldUsageVisitor(fpath, visited_fns, mjwarp_field_info) + for body in target_fn_node.body: + visitor.visit(body) + + render_context_in_caller = False + if len(target_fn_node.args.args) > 2: + third_param = target_fn_node.args.args[2] + if third_param.annotation: + annotation_str = ast.unparse(third_param.annotation) + render_context_in_caller = 'RenderContext' in annotation_str + + logging.info( + 'End trace function "%s". Output fields: %s, RenderContext: %s', + fn, visitor.data_out_fields, render_context_in_caller + ) + return FieldUsage( + model_fields=sorted(list(visitor.model_fields)), + data_fields=sorted(list(visitor.data_fields)), + data_out_fields=sorted(list(visitor.data_out_fields)), + render_context_in_caller=render_context_in_caller, + ) + + +def get_mjwarp_field_info(src: str, get_cls_type_annotations) -> Dict[str, FieldInfo]: + """Return field info for mujoco_warp/_src/types.py.""" + dataclass_map = { + 'opt': 'Option', + 'stat': 'Statistic', + 'efc': 'Constraint', + 'contact': 'Contact', + } + field_info = {} + type_classes = get_cls_type_annotations(src) + for field, typ in type_classes['Model'].items(): + if field == 'callback': + continue + if field in field_info: + raise AssertionError(f'Field {field} is duplicated in Model.') + if field in dataclass_map: + for sfield, styp in type_classes[dataclass_map[field]].items(): + field_name = field + '__' + sfield + field_info[field_name] = FieldInfo('Model', styp, (1, field_name)) + else: + field_info[field] = FieldInfo('Model', typ, (0, field)) + + for field, typ in type_classes['Data'].items(): + if field in field_info: + raise AssertionError(f'Field {field} is duplicated.') + if field in dataclass_map: + for sfield, styp in type_classes[dataclass_map[field]].items(): + field_name = field + '__' + sfield + field_info[field_name] = FieldInfo('Data', styp, (3, field_name)) + else: + field_info[field] = FieldInfo('Data', typ, (2, field)) + return field_info + + +def get_mjx_warp_field_info(src: str, get_cls_type_annotations) -> Dict[str, FieldInfo]: + """Return field info for mjx/warp/types.py.""" + field_info = {} + type_classes = get_cls_type_annotations(src) + for field, typ in type_classes['DataWarp'].items(): + field_info[field] = FieldInfo('Data', typ, (0, field)) + for field, typ in type_classes['ModelWarp'].items(): + field_info[field] = FieldInfo('Model', typ, (1, field)) + return field_info diff --git a/mjx/mujoco/mjx/codegen/update_for_mujoco_warp.sh b/mjx/mujoco/mjx/codegen/update_for_mujoco_warp.sh new file mode 100755 index 00000000..5186c6f5 --- /dev/null +++ b/mjx/mujoco/mjx/codegen/update_for_mujoco_warp.sh @@ -0,0 +1,119 @@ +#!/bin/bash +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +set -euo pipefail + +# --- Logging helpers --------------------------------------------------------- +log_stage() { echo -e "\n\033[1;34m==== $1 ====\033[0m"; } +log_ok() { echo -e "\033[1;32m ✓ $1\033[0m"; } +log_fail() { echo -e "\033[1;31m ✗ $1\033[0m"; } + +run_shim() { + local label="$1"; shift + log_stage "Generating shim: ${label}" + echo " → $*" + local output + if output=$("$@" --logtostderr 2>&1); then + # Show only Python-level log lines (filters noisy C++ infra logs). + echo "$output" | grep -E '\.py' || true + log_ok "${label}" + else + echo "$output" + log_fail "${label}" + exit 1 + fi +} + +# --- Path setup -------------------------------------------------------------- +mjwarp_base="mujoco/mjx/third_party/mujoco_warp/_src" +mjx_base="mujoco/mjx" + +# Derived paths (shared). +mjwarp="${mjwarp_base}" +mjx_warp_out="${mjx_base}/warp" +mjx_types="${mjx_base}/_src/types.py" +mjx_warp_types="${mjx_base}/warp/types.py" + +log_stage "Path configuration" +echo " mjwarp_base = ${mjwarp_base}" +echo " mjx_base = ${mjx_base}" +echo " mjx_types = ${mjx_types}" +echo " output dir = ${mjx_warp_out}" + +# --- Stage 1: Generate warp types ------------------------------------------- +log_stage "Stage 1/3: Generating warp types" +python mujoco/mjx/codegen/generate_warp_types.py \ + --mjx_warp_types_out_path=${mjx_warp_types} \ + --mjx_types_path=${mjx_types} +log_ok "Warp types written to ${mjx_warp_types}" + +# --- Stage 2: Build shim generator ------------------------------------------ +log_stage "Stage 2/3: Building shim generator" +generate_warp_shim="python mujoco/mjx/codegen/generate_warp_shim.py" +log_ok "Shim generator ready" + +# --- Stage 3: Generate shim code for each function -------------------------- +log_stage "Stage 3/3: Generating shim code" + +# Smooth. +run_shim "smooth:kinematics" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/smooth.py:kinematics \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py + +run_shim "smooth:tendon" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/smooth.py:tendon \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py \ + --append_to_output_file=True + +run_shim "smooth:com_pos" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/smooth.py:com_pos \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py \ + --append_to_output_file=True + +# Collision. +run_shim "collision_driver:collision" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/collision_driver.py:collision \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py + +# Forward. +run_shim "forward:forward" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/forward.py:forward \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py + +run_shim "forward:step" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/forward.py:step \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py \ + --append_to_output_file=True + +# Render and bvh. +run_shim "bvh:refit_bvh" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/bvh.py:refit_bvh \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py + +run_shim "render:render" ${generate_warp_shim} \ + --mjwarp_function=${mjwarp}/render.py:render \ + --mjx_warp_output_path=${mjx_warp_out}/ \ + --mjwarp_types=${mjwarp}/types.py + +log_stage "Done" +log_ok "All shims generated successfully" diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py index dbecb9c1..47b94ed2 100644 --- a/mjx/mujoco/mjx/warp/collision_driver.py +++ b/mjx/mujoco/mjx/warp/collision_driver.py @@ -46,7 +46,6 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) - @ffi.format_args_for_warp def _collision_shim( # Model diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py index 2f46f1f3..f260ea92 100644 --- a/mjx/mujoco/mjx/warp/forward.py +++ b/mjx/mujoco/mjx/warp/forward.py @@ -46,7 +46,6 @@ _cb = mjwp_types.Callback( **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init} ) - @ffi.format_args_for_warp def _forward_shim( # Model diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index 65b614b9..9d3e2f85 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -39,6 +39,10 @@ dependencies = [ warp = [ "warp-lang==1.12.1", ] +dev = [ + "isort", + "pyink", +] [project.scripts] mjx-testspeed = "mujoco.mjx.testspeed:main"