mjx

MuJoCo MJX dynamics adapters for OpenSCvx.

The recommended entry-point is MjxDynamics, a DynamicsAdapter that goes directly into the dynamics= slot of Problem and exposes the synthesized State/Control objects on .states / .controls::

dyn = ox.MjxDynamics(mjx_model)
problem = ox.Problem(
    dynamics=dyn,
    states=dyn.states,
    controls=dyn.controls,
    ...
)

Free-joint quaternion kinematics (nq > nv models such as drones or humanoids) are detected and handled automatically.

The lower-level mjx_dynamics callable factory is also public for advanced users who need to assemble their own BYOF dynamics dict (e.g. with custom State/Control names or interleaved with other states).

Note

Time dilation is handled automatically by the BYOF lowering pipeline; all functions return physical (un-dilated) quantities.

MjxDynamics

Bases: DynamicsAdapter

First-class MJX dynamics adapter for Problem.

Wraps a mujoco.mjx model so it can be passed directly to the dynamics= argument of Problem. The adapter constructs default qpos / qvel State objects and a ctrl Control matching the model's nq / nv / nu, exposes them via .states / .controls, and routes the MJX forward dynamics through the BYOF channel internally — without requiring the user to know about BYOF at all.

Example

Cartpole (nq == nv)::

mj_model = mujoco.MjModel.from_xml_path("cartpole.xml")
mj_model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_CONTACT
mjx_model = mjx.put_model(mj_model)

dyn = ox.MjxDynamics(mjx_model)
problem = ox.Problem(
    dynamics=dyn,
    states=dyn.states,
    controls=dyn.controls,
    ...
)

Quadrotor with a free joint (nq > nv) — quaternion kinematics are inserted automatically::

dyn = ox.MjxDynamics(mjx_model)  # nq=7, nv=6
problem = ox.Problem(
    dynamics=dyn, states=dyn.states, controls=dyn.controls, ...
)

Custom State/Control names or shapes are not supported here on purpose — the whole point of the adapter is "I don't want to think about names." Drop to the lower-level mjx_dynamics helper if you need that control — construct your own State/Control objects, pass them in, and assemble the BYOF dict yourself.

Supported joint structure

• Free (mjJNT_FREE), slide (mjJNT_SLIDE), and hinge (mjJNT_HINGE) joints only.
• If the model contains any free joints, they must all come before any slide/hinge joints in the MuJoCo layout.
• Ball joints (mjJNT_BALL) are explicitly refused — they share nq=4, nv=3 with free joints but use different kinematics, and would silently produce wrong dynamics.

Construction raises NotImplementedError if any of these conditions are violated; fall back to mjx_dynamics for those cases.

Auto-populated bounds

• qpos.min / qpos.max are read from mjx_model.jnt_range for slide / hinge joints flagged jnt_limited=True; free-joint slots and unlimited joints default to ±inf.
• ctrl.min / ctrl.max are read from actuator_ctrlrange for actuators flagged actuator_ctrllimited=True; otherwise ±inf.
• qvel bounds default to ±inf (MuJoCo has no per-joint velocity-limit concept).

Override any of these after construction if you want tighter problem-specific bounds.

Source code in openscvx/integrations/mjx.py

class MjxDynamics(DynamicsAdapter):
    """First-class MJX dynamics adapter for `Problem`.

    Wraps a ``mujoco.mjx`` model so it can be passed directly to the
    ``dynamics=`` argument of `Problem`. The adapter
    constructs default ``qpos`` / ``qvel`` State objects and a ``ctrl``
    Control matching the model's ``nq`` / ``nv`` / ``nu``, exposes them via
    ``.states`` / ``.controls``, and routes the MJX forward dynamics through
    the BYOF channel internally — without requiring the user to know about
    BYOF at all.

    Example:
        Cartpole (``nq == nv``)::

            mj_model = mujoco.MjModel.from_xml_path("cartpole.xml")
            mj_model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_CONTACT
            mjx_model = mjx.put_model(mj_model)

            dyn = ox.MjxDynamics(mjx_model)
            problem = ox.Problem(
                dynamics=dyn,
                states=dyn.states,
                controls=dyn.controls,
                ...
            )

        Quadrotor with a free joint (``nq > nv``) — quaternion kinematics
        are inserted automatically::

            dyn = ox.MjxDynamics(mjx_model)  # nq=7, nv=6
            problem = ox.Problem(
                dynamics=dyn, states=dyn.states, controls=dyn.controls, ...
            )

    Custom State/Control names or shapes are *not* supported here on purpose
    — the whole point of the adapter is "I don't want to think about names."
    Drop to the lower-level `mjx_dynamics` helper if you need that control —
    construct your own State/Control objects, pass them in, and assemble the
    BYOF dict yourself.

    Supported joint structure:
        * Free (``mjJNT_FREE``), slide (``mjJNT_SLIDE``), and hinge
          (``mjJNT_HINGE``) joints only.
        * If the model contains any free joints, they must all come
          *before* any slide/hinge joints in the MuJoCo layout.
        * Ball joints (``mjJNT_BALL``) are explicitly refused — they share
          ``nq=4, nv=3`` with free joints but use different kinematics, and
          would silently produce wrong dynamics.

        Construction raises ``NotImplementedError`` if any of these
        conditions are violated; fall back to `mjx_dynamics` for those
        cases.

    Auto-populated bounds:
        * ``qpos.min`` / ``qpos.max`` are read from ``mjx_model.jnt_range``
          for slide / hinge joints flagged ``jnt_limited=True``; free-joint
          slots and unlimited joints default to ``±inf``.
        * ``ctrl.min`` / ``ctrl.max`` are read from ``actuator_ctrlrange``
          for actuators flagged ``actuator_ctrllimited=True``; otherwise
          ``±inf``.
        * ``qvel`` bounds default to ``±inf`` (MuJoCo has no per-joint
          velocity-limit concept).

        Override any of these after construction if you want tighter
        problem-specific bounds.
    """

    def __init__(
        self,
        mjx_model: Any,
        *,
        return_component: str = "qacc",
        extra_postprocess: Optional[Callable[[Any], Any]] = None,
    ) -> None:
        from openscvx.symbolic.expr.control import Control
        from openscvx.symbolic.expr.state import State

        _validate_supported_joints(mjx_model)

        self.mjx_model = mjx_model
        self.return_component = return_component
        self.extra_postprocess = extra_postprocess

        nq = int(mjx_model.nq)
        nv = int(mjx_model.nv)
        nu = int(mjx_model.nu)

        self._qpos = State("qpos", shape=(nq,))
        self._qvel = State("qvel", shape=(nv,))
        self._ctrl = Control("ctrl", shape=(nu,))

        # Auto-populate bounds from the model so the user doesn't have to
        # re-type joint / actuator limits already declared in MJCF. Users
        # can still override any of these after construction.
        qpos_min, qpos_max, qvel_min, qvel_max, ctrl_min, ctrl_max = _initial_bounds_from_model(
            mjx_model, nq, nv, nu
        )
        self._qpos.min = qpos_min
        self._qpos.max = qpos_max
        self._qvel.min = qvel_min
        self._qvel.max = qvel_max
        self._ctrl.min = ctrl_min
        self._ctrl.max = ctrl_max

        self.states: list[State] = [self._qpos, self._qvel]
        self.controls: list[Control] = [self._ctrl]

    def expand(self) -> Tuple[dict, dict]:
        """Return ``(dynamics_dict, byof_dict)`` for this MJX model.

        - ``nq == nv``: ``dynamics_dict = {"qpos": qvel}`` (symbolic
          kinematic identity), ``byof_dict["dynamics"] = {"qvel": ...}``.
        - ``nq > nv``: ``dynamics_dict = {}``, ``byof_dict["dynamics"]``
          contains both ``"qpos"`` (quaternion kinematics) and ``"qvel"``.
        """
        nq = int(self.mjx_model.nq)
        nv = int(self.mjx_model.nv)

        byof_dynamics: dict = {
            "qvel": mjx_dynamics(
                self.mjx_model,
                qpos=self._qpos,
                qvel=self._qvel,
                ctrl=self._ctrl,
                return_component=self.return_component,
                extra_postprocess=self.extra_postprocess,
            ),
        }

        n_free = nq - nv
        if n_free > 0:
            byof_dynamics["qpos"] = _free_joint_qpos_dynamics(
                qpos=self._qpos,
                qvel=self._qvel,
                n_free_joints=n_free,
            )
            dynamics_dict: dict = {}
        else:
            dynamics_dict = {"qpos": self._qvel}

        return dynamics_dict, {"dynamics": byof_dynamics}

expand() -> Tuple[dict, dict]

Return (dynamics_dict, byof_dict) for this MJX model.

• nq == nv: dynamics_dict = {"qpos": qvel} (symbolic kinematic identity), byof_dict["dynamics"] = {"qvel": ...}.
• nq > nv: dynamics_dict = {}, byof_dict["dynamics"] contains both "qpos" (quaternion kinematics) and "qvel".

Source code in openscvx/integrations/mjx.py

def expand(self) -> Tuple[dict, dict]:
    """Return ``(dynamics_dict, byof_dict)`` for this MJX model.

    - ``nq == nv``: ``dynamics_dict = {"qpos": qvel}`` (symbolic
      kinematic identity), ``byof_dict["dynamics"] = {"qvel": ...}``.
    - ``nq > nv``: ``dynamics_dict = {}``, ``byof_dict["dynamics"]``
      contains both ``"qpos"`` (quaternion kinematics) and ``"qvel"``.
    """
    nq = int(self.mjx_model.nq)
    nv = int(self.mjx_model.nv)

    byof_dynamics: dict = {
        "qvel": mjx_dynamics(
            self.mjx_model,
            qpos=self._qpos,
            qvel=self._qvel,
            ctrl=self._ctrl,
            return_component=self.return_component,
            extra_postprocess=self.extra_postprocess,
        ),
    }

    n_free = nq - nv
    if n_free > 0:
        byof_dynamics["qpos"] = _free_joint_qpos_dynamics(
            qpos=self._qpos,
            qvel=self._qvel,
            n_free_joints=n_free,
        )
        dynamics_dict: dict = {}
    else:
        dynamics_dict = {"qpos": self._qvel}

    return dynamics_dict, {"dynamics": byof_dynamics}

mjx_dynamics(mjx_model: Any, *, qpos: State | slice, qvel: State | slice, ctrl: Control | slice, return_component: str = 'qacc', extra_postprocess: Optional[Callable[[Any], Any]] = None) -> Callable

Wrap a MuJoCo MJX model as a BYOF dynamics function.

Parameters:

Name	Type	Description	Default
mjx_model	Any	A model produced by :func:mujoco.mjx.put_model. Must be a JAX pytree (the standard MJX representation).	required
qpos	State | slice	Position state (or slice into the unified x vector). Length must equal mjx_model.nq.	required
qvel	State | slice	Velocity state (or slice). Length must equal mjx_model.nv.	required
ctrl	Control | slice	Control variable (or slice into the unified u vector). Length must equal mjx_model.nu.	required
return_component	str	Which MJX field to return. "qacc" (default) returns the generalized acceleration qacc for use as the qvel state derivative. "qvel" returns qvel for use as the qpos state derivative (rarely needed because that is already symbolic).	'qacc'
extra_postprocess	Optional[Callable[[Any], Any]]	Optional callable applied to the MJX data after mjx.forward. Useful for computing custom outputs (e.g. site positions) used elsewhere.	None

Returns:

Type	Description
Callable	A function f(x, u, node, params) -> jnp.ndarray matching the BYOF
Callable	dynamics signature.

Raises:

Type	Description
ImportError	If mujoco.mjx is not installed.
ValueError	If return_component is not one of the allowed values.

Note

MJX's contact solver uses lax.while_loop, which is not reverse-mode differentiable. For contact-free systems (manipulators, cartpoles, etc.) disable contacts before uploading the model::

mj_model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_CONTACT
mjx_model = mjx.put_model(mj_model)

Example

Cartpole swing-up dynamics::

import mujoco
import mujoco.mjx as mjx
import openscvx as ox
from openscvx.integrations import mjx_dynamics

mj_model = mujoco.MjModel.from_xml_path("cartpole.xml")
mjx_model = mjx.put_model(mj_model)

qpos = ox.State("qpos", shape=(mjx_model.nq,))
qvel = ox.State("qvel", shape=(mjx_model.nv,))
ctrl = ox.Control("ctrl", shape=(mjx_model.nu,))

qvel_dynamics = mjx_dynamics(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)

problem = ox.Problem(
    dynamics={"qpos": qvel},
    byof={"dynamics": {"qvel": qvel_dynamics}},
    states=[qpos, qvel],
    controls=[ctrl],
    ...
)

Source code in openscvx/integrations/mjx.py

def mjx_dynamics(
    mjx_model: Any,
    *,
    qpos: "State | slice",
    qvel: "State | slice",
    ctrl: "Control | slice",
    return_component: str = "qacc",
    extra_postprocess: Optional[Callable[[Any], Any]] = None,
) -> Callable:
    """Wrap a MuJoCo MJX model as a BYOF dynamics function.

    Args:
        mjx_model: A model produced by :func:`mujoco.mjx.put_model`. Must be a
            JAX pytree (the standard MJX representation).
        qpos: Position state (or slice into the unified ``x`` vector).
            Length must equal ``mjx_model.nq``.
        qvel: Velocity state (or slice). Length must equal ``mjx_model.nv``.
        ctrl: Control variable (or slice into the unified ``u`` vector).
            Length must equal ``mjx_model.nu``.
        return_component: Which MJX field to return. ``"qacc"`` (default)
            returns the generalized acceleration ``qacc`` for use as the
            ``qvel`` state derivative. ``"qvel"`` returns ``qvel`` for use as
            the ``qpos`` state derivative (rarely needed because that is
            already symbolic).
        extra_postprocess: Optional callable applied to the MJX ``data`` after
            ``mjx.forward``. Useful for computing custom outputs (e.g. site
            positions) used elsewhere.

    Returns:
        A function ``f(x, u, node, params) -> jnp.ndarray`` matching the BYOF
        dynamics signature.

    Raises:
        ImportError: If ``mujoco.mjx`` is not installed.
        ValueError: If ``return_component`` is not one of the allowed values.

    Note:
        MJX's contact solver uses ``lax.while_loop``, which is **not**
        reverse-mode differentiable. For contact-free systems (manipulators,
        cartpoles, etc.) disable contacts before uploading the model::

            mj_model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_CONTACT
            mjx_model = mjx.put_model(mj_model)

    Example:
        Cartpole swing-up dynamics::

            import mujoco
            import mujoco.mjx as mjx
            import openscvx as ox
            from openscvx.integrations import mjx_dynamics

            mj_model = mujoco.MjModel.from_xml_path("cartpole.xml")
            mjx_model = mjx.put_model(mj_model)

            qpos = ox.State("qpos", shape=(mjx_model.nq,))
            qvel = ox.State("qvel", shape=(mjx_model.nv,))
            ctrl = ox.Control("ctrl", shape=(mjx_model.nu,))

            qvel_dynamics = mjx_dynamics(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)

            problem = ox.Problem(
                dynamics={"qpos": qvel},
                byof={"dynamics": {"qvel": qvel_dynamics}},
                states=[qpos, qvel],
                controls=[ctrl],
                ...
            )
    """
    try:
        import mujoco.mjx as mjx
    except ImportError as e:
        raise ImportError(
            "mujoco.mjx is required for mjx_dynamics. Install with: pip install openscvx[mjx]"
        ) from e

    if return_component not in ("qacc", "qvel"):
        raise ValueError(f"return_component must be 'qacc' or 'qvel', got {return_component!r}")

    # Store the raw args; slices are resolved lazily on first call so that
    # mjx_dynamics() can be called before Problem construction assigns .slice.
    _qpos_arg = qpos
    _qvel_arg = qvel
    _ctrl_arg = ctrl
    _resolved: list = []  # populated on first call: [qpos_slice, qvel_slice, ctrl_slice]

    def f(x, u, node, params):
        del node, params  # MJX dynamics are stateless w.r.t. node and OpenSCvx params
        if not _resolved:
            _resolved.append(_resolve_slice(_qpos_arg, "qpos"))
            _resolved.append(_resolve_slice(_qvel_arg, "qvel"))
            _resolved.append(_resolve_slice(_ctrl_arg, "ctrl"))
        qpos_slice, qvel_slice, ctrl_slice = _resolved

        qpos_val = x[qpos_slice]
        qvel_val = x[qvel_slice]
        ctrl_val = u[ctrl_slice]

        data = mjx.make_data(mjx_model)
        data = data.replace(qpos=qpos_val, qvel=qvel_val, ctrl=ctrl_val)
        data = mjx.forward(mjx_model, data)

        if extra_postprocess is not None:
            data = extra_postprocess(data)

        if return_component == "qacc":
            return jnp.asarray(data.qacc)
        return jnp.asarray(data.qvel)

    return f