openscvx 2.dev3__py3-none-any.whl → 2.dev5__py3-none-any.whl

openscvx/__init__.py

@@ -25,6 +25,7 @@ from openscvx.discretization import (
     VectorizeDiscretizeLinearize,
 )
 from openscvx.expert import ByofSpec
+from openscvx.integrations import DynamicsAdapter, MjxDynamics
 from openscvx.loader import load_dict, load_json, load_yaml
 from openscvx.problem import Problem
 from openscvx.solvers import PTRSolver
@@ -176,6 +177,9 @@ __all__ = [
     "lie",
     # Expert mode types
     "ByofSpec",
+    # External-backend dynamics adapters
+    "DynamicsAdapter",
+    "MjxDynamics",
     # Discretization
     "DiscretizeLinearizeVectorize",
     "LinearizeDiscretize",

openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '2.dev3'
-__version_tuple__ = version_tuple = (2, 'dev3')
+__version__ = version = '2.dev5'
+__version_tuple__ = version_tuple = (2, 'dev5')
 
 __commit_id__ = commit_id = None

openscvx/integrations/__init__.py

@@ -1,47 +1,46 @@
-"""Adapters for MuJoCo MJX dynamics in OpenSCvx BYOF.
+"""External-backend dynamics adapters for OpenSCvx.
 
-The recommended entry-point is :func:`mjx_byof`, which returns a complete
-``byof["dynamics"]`` dict and automatically handles free-joint quaternion
-kinematics for floating-base models (drones, humanoids, etc.):
+The recommended entry-point is `MjxDynamics`, which goes directly into the
+``dynamics=`` slot of `Problem` and constructs the matching State/Control
+objects for the user. Free-joint quaternion kinematics for floating-base
+models (drones, humanoids) are detected and handled automatically::
 
-    from openscvx.integrations import mjx_byof
+    from openscvx.integrations import MjxDynamics
 
-    byof = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
+    dyn = MjxDynamics(mjx_model)
+    problem = ox.Problem(
+        dynamics=dyn,
+        states=dyn.states,
+        controls=dyn.controls,
+        ...
+    )
 
-For models without free joints (cartpoles, manipulators) the returned dict
-contains only ``"qvel"`` and ``dynamics={"qpos": qvel}`` should still be
-provided to :class:`~openscvx.Problem`.  For models with free joints
-(``nq > nv``) ``"qpos"`` is included automatically — no extra imports needed.
+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.
 
-:func:`mjx_dynamics` is also available for advanced users who need direct
-access to the BYOF callable for the ``qvel`` (acceleration) derivative.
+All MJX symbols delegate lazily so ``mujoco.mjx`` is only imported when
+actually used. The ``menagerie`` submodule is also loaded lazily.
 
-All symbols delegate lazily so ``mujoco.mjx`` is only imported when used.
-The :mod:`menagerie` submodule is loaded lazily via attribute access.
+Example — cartpole (``nq == nv``)::
 
-Example — cartpole (nq == nv)::
+    from openscvx.integrations import MjxDynamics
 
-    from openscvx.integrations import mjx_byof
+    dyn = MjxDynamics(mjx_model)
+    problem = ox.Problem(dynamics=dyn, states=dyn.states, controls=dyn.controls, ...)
 
-    byof = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-    problem = ox.Problem(dynamics={"qpos": qvel}, byof=byof, ...)
+Example — quadrotor with free joint (``nq=7``, ``nv=6``)::
 
-Example — quadrotor with free joint (nq=7, nv=6)::
+    from openscvx.integrations import MjxDynamics
 
-    from openscvx.integrations import mjx_byof
-
-    byof = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-    problem = ox.Problem(dynamics={}, byof=byof, ...)
+    dyn = MjxDynamics(mjx_model)
+    problem = ox.Problem(dynamics=dyn, states=dyn.states, controls=dyn.controls, ...)
 """
 
 from typing import Any
 
-
-def mjx_byof(*args: Any, **kwargs: Any) -> Any:
-    """Lazy delegate; imports ``mujoco.mjx`` on first call."""
-    from .mjx import mjx_byof as _mjx_byof
-
-    return _mjx_byof(*args, **kwargs)
+from .base import DynamicsAdapter
+from .mjx import MjxDynamics
 
 
 def mjx_dynamics(*args: Any, **kwargs: Any) -> Any:
@@ -59,4 +58,9 @@ def __getattr__(name: str) -> Any:
     raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
-__all__ = ["mjx_byof", "mjx_dynamics", "menagerie"]
+__all__ = [
+    "DynamicsAdapter",
+    "MjxDynamics",
+    "mjx_dynamics",
+    "menagerie",
+]

openscvx/integrations/base.py

@@ -0,0 +1,89 @@
+"""Base class for external-backend dynamics adapters.
+
+A `DynamicsAdapter` is the easy on-ramp for users who want to plug an
+external physics backend (MuJoCo MJX, Brax, Drake, ...) into OpenSCvx without
+manually constructing State/Control objects with matching shapes or routing
+raw JAX callables through the expert ``byof`` channel.
+
+The intended call site is::
+
+    dyn = ox.MjxDynamics(mjx_model)
+    problem = ox.Problem(
+        dynamics=dyn,
+        states=dyn.states,
+        controls=dyn.controls,
+        ...
+    )
+
+Internally, `Problem` detects the adapter, calls `DynamicsAdapter.expand`,
+and merges the resulting BYOF callables into the user's ``byof`` dict (if
+any). Everything downstream sees ordinary ``dynamics`` and ``byof`` dicts.
+"""
+
+from __future__ import annotations
+
+import copy
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Tuple
+
+if TYPE_CHECKING:
+    from openscvx.symbolic.expr.control import Control
+    from openscvx.symbolic.expr.state import State
+
+
+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.
+        """
+
+
+def _merge_byof(user_byof: dict | None, extra_byof: dict) -> dict:
+    """Merge an adapter-synthesized BYOF dict into a user-provided one.
+
+    Only the ``"dynamics"`` sub-dict is deep-merged; other keys are taken
+    verbatim from whichever side provides them. Raises ``ValueError`` on any
+    key collision under ``"dynamics"`` — a user passing both
+    ``dynamics=ox.MjxDynamics(...)`` and ``byof={"dynamics": {"qvel": ...}}``
+    almost certainly has a bug, and silent override would mask it.
+    """
+    if not user_byof:
+        return copy.copy(extra_byof)
+
+    merged = dict(user_byof)
+    extra_dyn = extra_byof.get("dynamics", {})
+    user_dyn = user_byof.get("dynamics", {})
+
+    if extra_dyn:
+        collisions = set(user_dyn) & set(extra_dyn)
+        if collisions:
+            raise ValueError(
+                "DynamicsAdapter produced byof['dynamics'] entries that "
+                f"collide with user-provided byof['dynamics']: {sorted(collisions)}. "
+                "Drop the duplicate keys from your byof dict, or drop the adapter "
+                "and assemble byof['dynamics'] manually for full control."
+            )
+        merged["dynamics"] = {**user_dyn, **extra_dyn}
+
+    return merged

openscvx/integrations/mjx.py

@@ -1,29 +1,35 @@
-"""MuJoCo MJX dynamics adapters for OpenSCvx BYOF.
+"""MuJoCo MJX dynamics adapters for OpenSCvx.
 
-The recommended entry-point is :func:`mjx_byof`, which returns a complete
-``byof["dynamics"]`` dict and automatically handles free-joint quaternion
-kinematics — no separate imports required:
+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``::
 
-    byof = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
+    dyn = ox.MjxDynamics(mjx_model)
+    problem = ox.Problem(
+        dynamics=dyn,
+        states=dyn.states,
+        controls=dyn.controls,
+        ...
+    )
 
-For models **without** free joints (cartpoles, manipulators, etc.) the
-returned dict contains only ``"qvel"``, and qpos kinematics must still be
-specified symbolically via ``dynamics={"qpos": qvel}``.  For models **with**
-free joints (drones, humanoids) ``"qpos"`` is included automatically and no
-symbolic dynamics entry is needed.
+Free-joint quaternion kinematics (``nq > nv`` models such as drones or
+humanoids) are detected and handled automatically.
 
-The lower-level :func:`mjx_dynamics` is also public for advanced users who
-need direct access to the BYOF callable for the ``qvel`` derivative.
+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.
 """
 
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple
 
 import jax.numpy as jnp
 
+from openscvx.integrations.base import DynamicsAdapter
+
 if TYPE_CHECKING:
     from openscvx.symbolic.expr.control import Control
     from openscvx.symbolic.expr.state import State
@@ -164,7 +170,7 @@ def _free_joint_qpos_dynamics(
 ) -> Callable:
     """BYOF callable for ``qpos`` when the model has quaternion free joints.
 
-    Used internally by :func:`mjx_byof`.  When a MuJoCo model has a
+    Used internally by `MjxDynamics`.  When a MuJoCo model has a
     floating-base free joint, ``nq > nv`` because each quaternion orientation
     contributes 4 position DOF but only 3 angular velocity DOF. The simple
     symbolic shorthand ``"qpos": qvel`` therefore fails a shape check. This
@@ -246,78 +252,245 @@ def _free_joint_qpos_dynamics(
     return f
 
 
-def mjx_byof(
-    mjx_model: Any,
-    *,
-    qpos: "State | slice",
-    qvel: "State | slice",
-    ctrl: "Control | slice",
-    return_component: str = "qacc",
-    extra_postprocess: Optional[Callable[[Any], Any]] = None,
-) -> dict:
-    """Return a complete ``byof["dynamics"]`` dict for a MuJoCo MJX model.
+# MuJoCo joint type enum (matches mujoco.mjtJoint): 0=free, 1=ball, 2=slide, 3=hinge.
+# Inlined here so we don't require `mujoco` to be importable just for type validation —
+# the user already needs mujoco to have constructed the mjx_model, but keeping the
+# numeric constants local makes this file self-contained.
+_MJ_JNT_FREE = 0
+_MJ_JNT_BALL = 1
+_MJ_JNT_SLIDE = 2
+_MJ_JNT_HINGE = 3
 
-    This is the recommended high-level entry-point.  It inspects the model's
-    ``nq`` and ``nv`` to detect free joints and automatically includes the
-    quaternion kinematics callable for ``qpos`` when needed.
 
-    Args:
-        mjx_model: A model produced by :func:`mujoco.mjx.put_model`.
-        qpos: Position state (or slice). Length must equal ``mjx_model.nq``.
-        qvel: Velocity state (or slice). Length must equal ``mjx_model.nv``.
-        ctrl: Control variable (or slice). Length must equal ``mjx_model.nu``.
-        return_component: Passed to :func:`mjx_dynamics`. ``"qacc"``
-            (default) uses the generalized acceleration as the ``qvel``
-            derivative; ``"qvel"`` returns qvel directly (rarely needed).
-        extra_postprocess: Optional callable applied to the MJX ``data``
-            object after ``mjx.forward``. Passed through to
-            :func:`mjx_dynamics`.
+def _initial_bounds_from_model(
+    mjx_model: Any, nq: int, nv: int, nu: int
+) -> Tuple[Any, Any, Any, Any, Any, Any]:
+    """Pull qpos / ctrl bounds out of the MJX model.
 
-    Returns:
-        A dict suitable for use as ``byof["dynamics"]``.
-        For models **without** free joints (``nq == nv``) only ``"qvel"`` is
-        included; position kinematics should still be provided symbolically
-        via ``dynamics={"qpos": qvel}``.
-        For models **with** free joints (``nq > nv``) both ``"qpos"`` and
-        ``"qvel"`` are included and no symbolic ``dynamics`` entry is needed.
+    Returns ``(qpos_min, qpos_max, qvel_min, qvel_max, ctrl_min, ctrl_max)``.
+
+    - ``qpos`` bounds come from ``mjx_model.jnt_range`` for slide/hinge joints
+      flagged ``jnt_limited=True``. All other qpos slots — free-joint
+      translations and quaternion components, unlimited slide/hinge joints —
+      get ``±inf``.
+    - ``ctrl`` bounds come from ``mjx_model.actuator_ctrlrange`` for actuators
+      flagged ``actuator_ctrllimited=True``. Unlimited actuators get ``±inf``.
+    - ``qvel`` bounds are always ``±inf`` because MuJoCo has no per-joint
+      velocity-limit concept; users override as needed.
+    """
+    import numpy as _np
+
+    qpos_min = _np.full(nq, -_np.inf)
+    qpos_max = _np.full(nq, _np.inf)
+    qvel_min = _np.full(nv, -_np.inf)
+    qvel_max = _np.full(nv, _np.inf)
+    ctrl_min = _np.full(nu, -_np.inf)
+    ctrl_max = _np.full(nu, _np.inf)
+
+    # Per-joint qpos bounds — only slide/hinge can be range-limited; free
+    # joints always have jnt_limited=False so we skip them safely.
+    jnt_type = _np.asarray(mjx_model.jnt_type).astype(int)
+    jnt_qposadr = _np.asarray(mjx_model.jnt_qposadr).astype(int)
+    jnt_limited = _np.asarray(mjx_model.jnt_limited).astype(bool)
+    jnt_range = _np.asarray(mjx_model.jnt_range).astype(float)
+    for i, jtype in enumerate(jnt_type):
+        if jtype in (_MJ_JNT_SLIDE, _MJ_JNT_HINGE) and jnt_limited[i]:
+            adr = int(jnt_qposadr[i])
+            qpos_min[adr] = jnt_range[i, 0]
+            qpos_max[adr] = jnt_range[i, 1]
+
+    if nu > 0:
+        act_limited = _np.asarray(mjx_model.actuator_ctrllimited).astype(bool)
+        act_range = _np.asarray(mjx_model.actuator_ctrlrange).astype(float)
+        for i in range(nu):
+            if act_limited[i]:
+                ctrl_min[i] = act_range[i, 0]
+                ctrl_max[i] = act_range[i, 1]
+
+    return qpos_min, qpos_max, qvel_min, qvel_max, ctrl_min, ctrl_max
+
+
+def _validate_supported_joints(mjx_model: Any) -> None:
+    """Refuse models whose joint layout the adapter cannot correctly handle.
+
+    `MjxDynamics` only supports models composed of free / slide / hinge joints
+    where all free joints precede the others in the state layout. Anything else
+    (ball joints, custom joint orderings) silently breaks the
+    `_free_joint_qpos_dynamics` arithmetic, so we refuse with a clear error
+    rather than producing wrong dynamics.
+    """
+    import numpy as _np
+
+    jnt_type = _np.asarray(mjx_model.jnt_type).astype(int)
+    supported = {_MJ_JNT_FREE, _MJ_JNT_SLIDE, _MJ_JNT_HINGE}
+    bad = sorted(set(jnt_type.tolist()) - supported)
+    if bad:
+        if _MJ_JNT_BALL in bad:
+            raise NotImplementedError(
+                "MjxDynamics does not support ball joints (mjJNT_BALL): they "
+                "share nq=4, nv=3 with free joints but use different "
+                "kinematics, and the current quaternion-kinematics callable "
+                "would silently produce wrong dynamics. Use `mjx_dynamics` "
+                "directly and assemble byof['dynamics'] manually."
+            )
+        raise NotImplementedError(
+            f"MjxDynamics only supports free, slide, and hinge joints; "
+            f"model contains unsupported joint types {bad}. Use "
+            "`mjx_dynamics` directly and assemble byof['dynamics'] manually."
+        )
+
+    # _free_joint_qpos_dynamics assumes all free joints come first in the
+    # state vector. If a slide/hinge precedes a free joint, the quaternion
+    # offsets would be off.
+    free_mask = jnt_type == _MJ_JNT_FREE
+    n_free = int(free_mask.sum())
+    if n_free and not free_mask[:n_free].all():
+        raise NotImplementedError(
+            "MjxDynamics requires all free joints to come before any "
+            "slide/hinge joints in the MuJoCo model. Reorder the joints in "
+            "your MJCF/URDF, or use `mjx_dynamics` directly to assemble "
+            "byof['dynamics'] yourself."
+        )
+
+
+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, no free joint)::
+        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)
 
-            byof = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
+            dyn = ox.MjxDynamics(mjx_model)
             problem = ox.Problem(
-                dynamics={"qpos": qvel},   # still required for non-free models
-                byof=byof, ...
+                dynamics=dyn,
+                states=dyn.states,
+                controls=dyn.controls,
+                ...
             )
 
-        Quadrotor / drone (nq > nv, one free joint)::
+        Quadrotor with a free joint (``nq > nv``) — quaternion kinematics
+        are inserted automatically::
 
-            byof = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
+            dyn = ox.MjxDynamics(mjx_model)  # nq=7, nv=6
             problem = ox.Problem(
-                dynamics={},               # qpos handled automatically
-                byof=byof, ...
+                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.
     """
-    nq = int(mjx_model.nq)
-    nv = int(mjx_model.nv)
-
-    result: dict = {
-        "qvel": mjx_dynamics(
-            mjx_model,
-            qpos=qpos,
-            qvel=qvel,
-            ctrl=ctrl,
-            return_component=return_component,
-            extra_postprocess=extra_postprocess,
-        ),
-    }
-
-    n_free = nq - nv  # each free joint contributes exactly 1 extra position DOF
-    if n_free > 0:
-        result["qpos"] = _free_joint_qpos_dynamics(
-            qpos=qpos,
-            qvel=qvel,
-            n_free_joints=n_free,
+
+    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 result
+        return dynamics_dict, {"dynamics": byof_dynamics}

openscvx/problem.py

@@ -42,6 +42,7 @@ from openscvx.discretization import (
     resolve_discretizer_config,
 )
 from openscvx.expert import ByofSpec
+from openscvx.integrations.base import DynamicsAdapter, _merge_byof
 from openscvx.lowered import LoweredProblem, ParameterDict
 from openscvx.lowered.dynamics import Dynamics
 from openscvx.lowered.jax_constraints import (
@@ -70,7 +71,7 @@ from openscvx.utils.caching import (
 class Problem:
     def __init__(
         self,
-        dynamics: dict,
+        dynamics: Union[dict, DynamicsAdapter],
         constraints: List[Union[Constraint, CTCS]],
         states: List[State],
         controls: List[Control],
@@ -92,9 +93,11 @@ class Problem:
         """The primary class in charge of compiling and exporting the solvers.
 
         Args:
-            dynamics (dict): Dictionary mapping state names to their dynamics expressions.
-                Each key should be a state name, and each value should be an Expr
-                representing the derivative of that state.
+            dynamics: Dictionary mapping state names to their dynamics
+                expressions. Each key should be a state name, and each value
+                should be an ``Expr`` representing the derivative of that
+                state. A ``DynamicsAdapter`` may also be passed here in
+                place of the dict.
             constraints (List[Union[CTCSConstraint, NodalConstraint]]):
                 List of constraints decorated with @ctcs or @nodal
             states (List[State]): List of State objects representing the state variables.
@@ -250,6 +253,12 @@ class Problem:
         self._float_dtype: str = float_dtype
 
         # Symbolic Preprocessing & Augmentation
+        # If `dynamics` is a DynamicsAdapter, expand it into the standard
+        # (dynamics_dict, byof_dict) representation and merge into user byof.
+        if isinstance(dynamics, DynamicsAdapter):
+            dynamics, adapter_byof = dynamics.expand()
+            byof = _merge_byof(byof, adapter_byof)
+
         # Resolve byof: dict → ByofSpec (validates keys and nested specs)
         if byof is not None:
             byof = ByofSpec.model_validate(byof)

{openscvx-2.dev3.dist-info → openscvx-2.dev5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 2.dev3
+Version: 2.dev5
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

{openscvx-2.dev3.dist-info → openscvx-2.dev5.dist-info}/RECORD

@@ -1,9 +1,9 @@
-openscvx/__init__.py,sha256=S-pt3p_BI5aOdOYPLrxKOshOeLEBXmqPoqZfROqq730,3542
+openscvx/__init__.py,sha256=xmpOTuaULy4YJFhqUpXGSSMu-vjTBJnFIhohXdi8ttc,3688
 openscvx/__main__.py,sha256=Hwm7mtVg3tLdvoUPkpcQv8KF3wxl72PNLBp9axFu8GY,2991
-openscvx/_version.py,sha256=hg3VdMRRozfeXtGifcZ3onI-rf5Dj_BRtApPZ_S-IvU,523
+openscvx/_version.py,sha256=weTnVwzOsaNm5p3YrseycYmlTphq5sEq-JrqG3LXouM,523
 openscvx/config.py,sha256=qfDDYoCe6WqJglKsx5b2W48YOglXenKr-PVRRdCFhYE,9898
 openscvx/loader.py,sha256=FvKLkkXd4ihd5FqLFF8Cd9VnPbPwTV_azBRnEipi28c,7654
-openscvx/problem.py,sha256=g-DKsdpZ5rOxEoMVOUj_QwB1L1edN0rgtRn8Z-TjgTI,47177
+openscvx/problem.py,sha256=GlCxYFhpacIt6Sxte3-K-d_GbGxVL9F9Qzy4y8-hrGc,47675
 openscvx/algorithms/__init__.py,sha256=f5VhjFb40JyVPgJJh6fUxRbmyQIglUyiNYL22nOMrgs,5632
 openscvx/algorithms/augmented_lagrangian.py,sha256=liEHtqONbpmw7CZJJCkAluPOEbFslFsvCElWHusHpn4,15807
 openscvx/algorithms/base.py,sha256=JsaVfS-hyHeGU1GUuQV9i-EpKQ1HBE0RHABIBGrEonM,30205
@@ -27,9 +27,10 @@ openscvx/expert/validation.py,sha256=ofOwg6t3LcrE-xaefBnEf0EAkUd9b5EwMQOjJYEYKLA
 openscvx/init/__init__.py,sha256=1nOhjqVgZRDTsHfozWPagPcyp99hskI3u31PF5kBzvw,893
 openscvx/init/interpolation.py,sha256=khypZhwADcYIhudn6EnLWMka9dmjY44aRpnGXhULQ3k,11165
 openscvx/init/inverse_kinematics.py,sha256=9mdBADa2TWvxi5z_wK_rmUdjk_K_3KJc5w2O9YfACDg,9084
-openscvx/integrations/__init__.py,sha256=sIvGvKbIrvDenhW31RohtiPzWtbFGN4vMgGTVXFuYG0,2132
+openscvx/integrations/__init__.py,sha256=UWMr7SHotwhUa9a636rxWPvpZJpeQOinHkLpQ9Apz9M,1959
+openscvx/integrations/base.py,sha256=oHDOPB-hTa9GFZ1tcmkdmzxG7pOkz4R5dT0XVGe5rKA,3377
 openscvx/integrations/menagerie.py,sha256=Zm2aGwwkqJPGQXjnHFoKI3dkXHhVzUB1eUIZUywmT28,6431
-openscvx/integrations/mjx.py,sha256=QQqygZMMzsuBexk8H93b8NWZy9fEDSR8ga6anhqecH0,11937
+openscvx/integrations/mjx.py,sha256=E71AOor70hR_wC7XGyACOGpSBaQklYVc6-RCqZghT0U,19112
 openscvx/integrators/__init__.py,sha256=easV2-ruQLif5e8UBqE-xP5mpslYCN84gzT3OfZGQgo,1616
 openscvx/integrators/diffrax.py,sha256=5RWNtaAUSGAWeHvX897fBxWjjJLVM4ORhqC12gdFVIA,4215
 openscvx/integrators/runge_kutta.py,sha256=yMf_JLVToPgdwayFP7ZZ-w-grf_ch9UgyPsZu3mYRb0,2972
@@ -137,9 +138,9 @@ openscvx/utils/caching.py,sha256=BPkT_IbmYT-i-BZ-himdWUc_4oBcwXWJxeUMwQWnSNc,934
 openscvx/utils/printing.py,sha256=zl3IxnKhwITqB5dK0Ru2IlORPqB61Y3DgsTryqMNu9M,13360
 openscvx/utils/profiling.py,sha256=k2x-i0CpG_kRe6dNcNBGu-ylrOtQw4B4C1UaOTjUMfU,1678
 openscvx/utils/utils.py,sha256=M25RHE_7DSr3Reaca0xCXnDSY9KHuqYvXdh5m1ZotEc,3047
-openscvx-2.dev3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-openscvx-2.dev3.dist-info/METADATA,sha256=h-MjAPo1a0EyE9EG_uW8hgBjBsyYWpBf5tPeRcOEpdQ,10662
-openscvx-2.dev3.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-openscvx-2.dev3.dist-info/entry_points.txt,sha256=1Oqek8Sy28hmAZFgZXDxFXYVf56YLYWlHjhh9RYJ7wE,52
-openscvx-2.dev3.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
-openscvx-2.dev3.dist-info/RECORD,,
+openscvx-2.dev5.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+openscvx-2.dev5.dist-info/METADATA,sha256=-i1S5BQH7E22WfWRdoGoe1dUnRoWiwt9DK8YfVVPAp4,10662
+openscvx-2.dev5.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+openscvx-2.dev5.dist-info/entry_points.txt,sha256=1Oqek8Sy28hmAZFgZXDxFXYVf56YLYWlHjhh9RYJ7wE,52
+openscvx-2.dev5.dist-info/top_level.txt,sha256=nUT4Ybefzh40H8tVXqc1RzKESy_MAowElb-CIvAbd4Q,9
+openscvx-2.dev5.dist-info/RECORD,,