integrations

External-backend dynamics adapters for OpenSCvx.

Two backends are currently supported. MjxDynamics wraps a MuJoCo MJX model; FraxDynamics wraps a frax.Robot (Fast Robot Kinematics and Dynamics in JAX). Both go directly into the dynamics= slot of Problem and construct the matching State/Control objects for the user::

from openscvx.integrations import FraxDynamics
from frax.robots.franka_panda import load_panda

dyn = FraxDynamics(load_panda())
problem = ox.Problem(dynamics=dyn, states=dyn.states, controls=dyn.controls, ...)

For MJX, free-joint quaternion kinematics for floating-base models (drones, humanoids) are detected and handled automatically::

from openscvx.integrations import MjxDynamics

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

For advanced users who need custom State/Control names (or to interleave them with extra custom states), mjx_dynamics is exposed as the underlying BYOF callable factory — assemble your own byof["dynamics"] dict from it.

All MJX symbols delegate lazily so mujoco.mjx is only imported when actually used. The menagerie submodule is also loaded lazily.

Example — cartpole (nq == nv)::

from openscvx.integrations import MjxDynamics

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

Example — quadrotor with free joint (nq=7, nv=6)::

from openscvx.integrations import MjxDynamics

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

DynamicsAdapter

Bases: ABC

Abstract base class for external-backend dynamics adapters.

Subclasses describe the State/Control objects they synthesize on .states / .controls and implement expand to return the two-channel (dynamics_dict, byof_dict) representation consumed by Problem.

The split mirrors the existing two-channel API: dynamics_dict carries symbolic Expr entries (e.g. {"qpos": qvel}) while byof_dict carries raw JAX callables under the "dynamics" key. Either or both may be empty, but expand() should never silently produce overlapping keys.

Source code in openscvx/integrations/base.py

class DynamicsAdapter(ABC):
    """Abstract base class for external-backend dynamics adapters.

    Subclasses describe the State/Control objects they synthesize on
    ``.states`` / ``.controls`` and implement `expand` to return the
    two-channel ``(dynamics_dict, byof_dict)`` representation consumed by
    `Problem`.

    The split mirrors the existing two-channel API: ``dynamics_dict`` carries
    symbolic Expr entries (e.g. ``{"qpos": qvel}``) while ``byof_dict`` carries
    raw JAX callables under the ``"dynamics"`` key. Either or both may be
    empty, but ``expand()`` should never silently produce overlapping keys.
    """

    states: list["State"]
    controls: list["Control"]

    @abstractmethod
    def expand(self) -> Tuple[dict, dict]:
        """Return ``(dynamics_dict, byof_dict)`` in OpenSCvx's internal form.

        ``dynamics_dict`` maps state names to symbolic ``Expr`` derivatives
        (the same shape as the ``dynamics=`` argument to ``Problem``).
        ``byof_dict`` has the same shape as the ``byof=`` argument: its
        ``"dynamics"`` key (if present) maps state names to raw JAX callables.
        """

expand() -> Tuple[dict, dict] abstractmethod

Return (dynamics_dict, byof_dict) in OpenSCvx's internal form.

dynamics_dict maps state names to symbolic Expr derivatives (the same shape as the dynamics= argument to Problem). byof_dict has the same shape as the byof= argument: its "dynamics" key (if present) maps state names to raw JAX callables.

Source code in openscvx/integrations/base.py

@abstractmethod
def expand(self) -> Tuple[dict, dict]:
    """Return ``(dynamics_dict, byof_dict)`` in OpenSCvx's internal form.

    ``dynamics_dict`` maps state names to symbolic ``Expr`` derivatives
    (the same shape as the ``dynamics=`` argument to ``Problem``).
    ``byof_dict`` has the same shape as the ``byof=`` argument: its
    ``"dynamics"`` key (if present) maps state names to raw JAX callables.
    """

FraxDynamics

Bases: DynamicsAdapter

First-class frax dynamics adapter for Problem.

Wraps a frax.Robot so it can be passed directly to the dynamics= argument of Problem. The adapter constructs default q / qd State objects and a tau Control matching the robot's joint count, exposes them via .states / .controls, and routes the frax forward dynamics through the BYOF channel internally — without requiring the user to know about BYOF.

Joint limits and torque bounds are read from the robot's URDF and auto-populated on the State/Control objects. Override any of them after construction if you want tighter problem-specific bounds::

dyn = ox.FraxDynamics(load_panda())
q, qd = dyn.states
q.initial = np.array([0.0, -0.7854, 0.0, -2.3562, 0.0, 1.5708, 0.7854])
# tighten velocity limits for a slow approach
qd.max = np.full(7, 1.0)

Example

Franka Panda::

from frax.robots.franka_panda import load_panda
import openscvx as ox

dyn = ox.FraxDynamics(load_panda())
problem = ox.Problem(
    dynamics=dyn,
    states=dyn.states,
    controls=dyn.controls,
    constraints=[...],
    N=50,
    time=ox.Time(...),
)

Floating-base robots (robot.includes_floating_dof = True) are supported. The tau Control covers only the num_actuated_joints actuated DOFs; the adapter inserts the required zero-torque entries for the unactuated floating-base DOFs internally.

Custom State/Control names are not supported — drop to the lower-level frax_dynamics helper for that.

Auto-populated bounds

• q.min / q.max: from the URDF <limit lower/upper>.
• qd.min / qd.max: from the URDF <limit velocity> (mirrored as ±max_vel).
• tau.min / tau.max: from the URDF <limit effort> for the actuated joints.

frax's ±1e6 "no limit" sentinel (used for floating-base DOFs) is mapped to ±inf on all of the above.

Source code in openscvx/integrations/frax.py

class FraxDynamics(DynamicsAdapter):
    """First-class frax dynamics adapter for `Problem`.

    Wraps a ``frax.Robot`` so it can be passed directly to the ``dynamics=``
    argument of `Problem`. The adapter constructs default ``q`` / ``qd`` State
    objects and a ``tau`` Control matching the robot's joint count, exposes
    them via ``.states`` / ``.controls``, and routes the frax forward dynamics
    through the BYOF channel internally — without requiring the user to know
    about BYOF.

    Joint limits and torque bounds are read from the robot's URDF and
    auto-populated on the State/Control objects. Override any of them after
    construction if you want tighter problem-specific bounds::

        dyn = ox.FraxDynamics(load_panda())
        q, qd = dyn.states
        q.initial = np.array([0.0, -0.7854, 0.0, -2.3562, 0.0, 1.5708, 0.7854])
        # tighten velocity limits for a slow approach
        qd.max = np.full(7, 1.0)

    Example:
        Franka Panda::

            from frax.robots.franka_panda import load_panda
            import openscvx as ox

            dyn = ox.FraxDynamics(load_panda())
            problem = ox.Problem(
                dynamics=dyn,
                states=dyn.states,
                controls=dyn.controls,
                constraints=[...],
                N=50,
                time=ox.Time(...),
            )

    Floating-base robots (``robot.includes_floating_dof = True``) are
    supported. The ``tau`` Control covers only the ``num_actuated_joints``
    actuated DOFs; the adapter inserts the required zero-torque entries for the
    unactuated floating-base DOFs internally.

    Custom State/Control names are *not* supported — drop to the lower-level
    `frax_dynamics` helper for that.

    Auto-populated bounds:
        * ``q.min`` / ``q.max``: from the URDF ``<limit lower/upper>``.
        * ``qd.min`` / ``qd.max``: from the URDF ``<limit velocity>``
          (mirrored as ``±max_vel``).
        * ``tau.min`` / ``tau.max``: from the URDF ``<limit effort>`` for the
          actuated joints.

        frax's ``±1e6`` "no limit" sentinel (used for floating-base DOFs) is
        mapped to ``±inf`` on all of the above.
    """

    def __init__(self, robot: Any) -> None:
        from openscvx.symbolic.expr.control import Control
        from openscvx.symbolic.expr.state import State

        self._robot = robot
        nj = robot.num_joints
        na = robot.num_actuated_joints

        self._q = State("q", shape=(nj,))
        self._qd = State("qd", shape=(nj,))
        self._tau = Control("tau", shape=(na,))

        q_min, q_max, qd_min, qd_max, tau_min, tau_max = _bounds_from_robot(robot)
        self._q.min = q_min
        self._q.max = q_max
        self._qd.min = qd_min
        self._qd.max = qd_max
        self._tau.min = tau_min
        self._tau.max = tau_max

        self.states: list[State] = [self._q, self._qd]
        self.controls: list[Control] = [self._tau]

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

        frax always has ``nq == nv`` (no quaternion floating-base split), so
        ``dynamics_dict`` is always ``{"q": qd}`` and the full joint
        acceleration goes through BYOF.
        """
        dynamics_dict = {"q": self._qd}
        byof_dict = {
            "dynamics": {
                "qd": frax_dynamics(
                    self._robot,
                    q=self._q,
                    qd=self._qd,
                    tau=self._tau,
                )
            }
        }
        return dynamics_dict, byof_dict

expand() -> Tuple[dict, dict]

Return (dynamics_dict, byof_dict) for this frax robot.

frax always has nq == nv (no quaternion floating-base split), so dynamics_dict is always {"q": qd} and the full joint acceleration goes through BYOF.

Source code in openscvx/integrations/frax.py

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

    frax always has ``nq == nv`` (no quaternion floating-base split), so
    ``dynamics_dict`` is always ``{"q": qd}`` and the full joint
    acceleration goes through BYOF.
    """
    dynamics_dict = {"q": self._qd}
    byof_dict = {
        "dynamics": {
            "qd": frax_dynamics(
                self._robot,
                q=self._q,
                qd=self._qd,
                tau=self._tau,
            )
        }
    }
    return dynamics_dict, byof_dict

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}

frax_dynamics(*args: Any, **kwargs: Any) -> Any

Lazy delegate to the low-level frax BYOF factory.

Source code in openscvx/integrations/__init__.py

def frax_dynamics(*args: Any, **kwargs: Any) -> Any:
    """Lazy delegate to the low-level frax BYOF factory."""
    from .frax import frax_dynamics as _frax_dynamics

    return _frax_dynamics(*args, **kwargs)

mjx_dynamics(*args: Any, **kwargs: Any) -> Any

Lazy delegate; imports mujoco.mjx on first call.

Source code in openscvx/integrations/__init__.py

def mjx_dynamics(*args: Any, **kwargs: Any) -> Any:
    """Lazy delegate; imports ``mujoco.mjx`` on first call."""
    from .mjx import mjx_dynamics as _mjx_dynamics

    return _mjx_dynamics(*args, **kwargs)