openscvx 2.dev4__tar.gz → 2.dev5__tar.gz

{openscvx-2.dev4/openscvx.egg-info → openscvx-2.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 2.dev4
+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.dev4 → openscvx-2.dev5}/examples/mjx/cartpole_mjx.py

@@ -34,7 +34,6 @@ except ImportError:
     sys.exit(1)
 
 import openscvx as ox
-from openscvx.integrations import mjx_byof
 
 CARTPOLE_XML = """
 <mujoco model="cartpole">
@@ -69,36 +68,32 @@ n_u = int(mjx_model.nu)
 n = 60
 total_time = 3.0
 
-qpos = ox.State("qpos", shape=(n_q,))
+# ── MJX dynamics as a first-class adapter ─────────────────────────────────────
+# `MjxDynamics` builds default qpos / qvel / ctrl State and Control objects
+# matching the model's nq / nv / nu and routes the MJX forward dynamics into
+# the BYOF channel internally — no separate `byof=` plumbing required.
+dyn = ox.MjxDynamics(mjx_model)
+qpos, qvel = dyn.states
+(ctrl,) = dyn.controls
+
 qpos.min = np.array([-3.0, -2.0 * np.pi])
 qpos.max = np.array([3.0, 2.0 * np.pi])
 qpos.initial = np.array([0.0, np.pi])
 qpos.final = np.array([0.0, 0.0])
 
-qvel = ox.State("qvel", shape=(n_v,))
 qvel.min = np.array([-10.0, -15.0])
 qvel.max = np.array([10.0, 15.0])
 qvel.initial = np.array([0.0, 0.0])
 qvel.final = np.array([0.0, 0.0])
 
-ctrl = ox.Control("ctrl", shape=(n_u,))
 ctrl.min = np.array([-1.0])
 ctrl.max = np.array([1.0])
 ctrl.guess = np.zeros((n, n_u))
 
-states = [qpos, qvel]
-controls = [ctrl]
-
-dynamics = {
-    "qpos": qvel,
-}
-
-byof: ox.ByofSpec = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-
 constraints = []
-for state in states:
+for state in dyn.states:
     constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
-for control in controls:
+for control in dyn.controls:
     constraints.extend([ox.ctcs(control <= control.max), ox.ctcs(control.min <= control)])
 
 theta_guess = np.linspace(np.pi, 0.0, n)
@@ -115,13 +110,12 @@ time = ox.Time(
 )
 
 problem = ox.Problem(
-    dynamics=dynamics,
-    states=states,
-    controls=controls,
+    dynamics=dyn,
+    states=dyn.states,
+    controls=dyn.controls,
     time=time,
     constraints=constraints,
     N=n,
-    byof=byof,
     algorithm={
         "lam_prox": 1e-1,
         "lam_cost": 1e-2,

{openscvx-2.dev4 → openscvx-2.dev5}/examples/mjx/double_cartpole_mjx.py

@@ -34,8 +34,6 @@ except ImportError:
     sys.exit(1)
 
 import openscvx as ox
-from openscvx import ByofSpec
-from openscvx.integrations import mjx_byof
 
 L1, L2 = 0.5, 0.4  # link lengths (m)
 
@@ -81,32 +79,26 @@ n_u = int(mjx_model.nu)  # 1: cart force
 n = 400
 total_time = 2.5
 
-# ── State / control definitions ───────────────────────────────────────────────
-qpos = ox.State("qpos", shape=(n_q,))
+# ── MJX dynamics adapter ──────────────────────────────────────────────────────
+dyn = ox.MjxDynamics(mjx_model)
+qpos, qvel = dyn.states
+(ctrl,) = dyn.controls
+ctrl.parameterization = "ZOH"
+
 qpos.min = np.array([-8.0, -2 * np.pi, -2 * np.pi])
 qpos.max = np.array([8.0, 2 * np.pi, 2 * np.pi])
 qpos.initial = np.array([0.0, np.pi, 0.0])  # cart at origin, link1 hanging down
 qpos.final = [ox.Free(0.0), 0.0, 0.0]  # cart free, both links upright
 
-qvel = ox.State("qvel", shape=(n_v,))
 qvel.min = np.array([-12.0, -12.0, -12.0])
 qvel.max = np.array([12.0, 12.0, 12.0])
 qvel.initial = np.zeros(n_v)
 qvel.final = [0.0, 0.0, 0.0]
 
-ctrl = ox.Control("ctrl", shape=(n_u,), parameterization="ZOH")
 ctrl.min = np.array([-2.0])
 ctrl.max = np.array([2.0])
 ctrl.guess = np.zeros((n, n_u))
 
-states = [qpos, qvel]
-controls = [ctrl]
-
-# ── Dynamics: position kinematics symbolically, velocity via MJX ──────────────
-dynamics: dict = {"qpos": qvel}  # nq == nv, valid for all-revolute/prismatic joints
-
-byof: ByofSpec = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-
 # ── Constraints (CTCS on state / control bounds) ───────────────────────────────
 constraints = []
 
@@ -123,13 +115,12 @@ time = ox.Time(
 )
 
 problem = ox.Problem(
-    dynamics=dynamics,
-    states=states,
-    controls=controls,
+    dynamics=dyn,
+    states=dyn.states,
+    controls=dyn.controls,
     time=time,
     constraints=constraints,
     N=n,
-    byof=byof,
     algorithm={
         "lam_prox": 1e-1,
         "lam_cost": 0e0,

{openscvx-2.dev4 → openscvx-2.dev5}/examples/mjx/skydio_x2_mjx.py

@@ -51,8 +51,7 @@ from examples.plotting_viser import (
     create_animated_plotting_server,
     create_scp_animated_plotting_server,
 )
-from openscvx import ByofSpec, Problem
-from openscvx.integrations import mjx_byof
+from openscvx import Problem
 from openscvx.utils import gen_vertices, rot
 
 HOVER_CTRL = 3.2495625  # N per motor for level hover (from menagerie keyframe)
@@ -78,33 +77,28 @@ n_u = int(mjx_model.nu)  # 4 — rotor thrusts
 n = 22
 total_time = 24.0
 
-# ── State / control definitions ───────────────────────────────────────────────
-qpos = ox.State("qpos", shape=(n_q,))
+# ── MJX dynamics adapter ──────────────────────────────────────────────────────
+# The free joint has nq=7 but nv=6 (quaternion adds one extra position DOF).
+# MjxDynamics detects nq > nv and automatically routes quaternion kinematics
+# for "qpos" alongside the MJX "qvel" dynamics through the BYOF channel.
+dyn = ox.MjxDynamics(mjx_model)
+qpos, qvel = dyn.states
+(ctrl,) = dyn.controls
+
 qpos.min = np.array([-200.0, -100.0, 15.0, -1.0, -1.0, -1.0, -1.0])
 qpos.max = np.array([200.0, 100.0, 200.0, 1.0, 1.0, 1.0, 1.0])
 qpos.initial = np.concatenate([START_POS, HOVER_QUAT])
 qpos.final = [10.0, 0.0, 20.0, ("free", 1.0), ("free", 0.0), ("free", 0.0), ("free", 0.0)]
 
-qvel = ox.State("qvel", shape=(n_v,))
 qvel.min = np.array([-100.0, -100.0, -100.0, -10.0, -10.0, -10.0])
 qvel.max = np.array([100.0, 100.0, 100.0, 10.0, 10.0, 10.0])
 qvel.initial = np.zeros(n_v)
 qvel.final = [("free", 0.0)] * n_v
 
-ctrl = ox.Control("ctrl", shape=(n_u,))
 ctrl.min = np.zeros(n_u)
 ctrl.max = 13.0 * np.ones(n_u)
 ctrl.guess = HOVER_CTRL * np.ones((n, n_u))
 
-states = [qpos, qvel]
-controls = [ctrl]
-
-# ── Dynamics via BYOF ─────────────────────────────────────────────────────────
-# The free joint has nq=7 but nv=6 (quaternion adds one extra position DOF).
-# nq=7, nv=6 (free joint): mjx_byof detects nq > nv and automatically
-# includes quaternion kinematics for "qpos" alongside the MJX "qvel" dynamics.
-byof: ByofSpec = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-
 # ── Gate parameters (matching examples/drone/drone_racing.py) ───────────────
 n_gates = 10
 initial_gate_centers = [
@@ -135,9 +129,9 @@ gate_centers = np.array(modified_centers)
 
 # ── Constraints ───────────────────────────────────────────────────────────────
 constraints = []
-for state in states:
+for state in dyn.states:
     constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
-for control in controls:
+for control in dyn.controls:
     constraints.extend([ox.ctcs(control <= control.max), ox.ctcs(control.min <= control)])
 
 # Enforce sequential gate traversal using nodal constraints on qpos position.
@@ -172,13 +166,12 @@ time = ox.Time(
 )
 
 problem = Problem(
-    dynamics={},  # all dynamics go through BYOF
-    states=states,
-    controls=controls,
+    dynamics=dyn,
+    states=dyn.states,
+    controls=dyn.controls,
     time=time,
     constraints=constraints,
     N=n,
-    byof=byof,
     algorithm={
         "lam_prox": 1e-1,
         "lam_cost": 1e-2,

{openscvx-2.dev4 → openscvx-2.dev5}/examples/mjx/triple_cartpole_3d_mjx.py

@@ -57,8 +57,7 @@ except ImportError:
     sys.exit(1)
 
 import openscvx as ox
-from openscvx import ByofSpec, Problem
-from openscvx.integrations import mjx_byof
+from openscvx import Problem
 
 L1, L2, L3 = 0.5, 0.4, 0.3  # link lengths (m)
 
@@ -122,8 +121,11 @@ IDX_A1, IDX_B1 = 2, 3
 IDX_A2, IDX_B2 = 4, 5
 IDX_A3, IDX_B3 = 6, 7
 
-# ── State / control definitions ───────────────────────────────────────────────
-qpos = ox.State("qpos", shape=(n_q,))
+# ── MJX dynamics adapter ──────────────────────────────────────────────────────
+dyn = ox.MjxDynamics(mjx_model)
+qpos, qvel = dyn.states
+(ctrl,) = dyn.controls
+
 qpos.min = np.array(
     [-8.0, -8.0, -2 * np.pi, -2 * np.pi, -2 * np.pi, -2 * np.pi, -2 * np.pi, -2 * np.pi]
 )
@@ -142,25 +144,15 @@ qpos.final = [
     0.0,  # link 3 upright
 ]
 
-qvel = ox.State("qvel", shape=(n_v,))
 qvel.min = np.full(n_v, -12.0)
 qvel.max = np.full(n_v, 12.0)
 qvel.initial = np.zeros(n_v)
 qvel.final = [0.0] * n_v
 
-ctrl = ox.Control("ctrl", shape=(n_u,))
 ctrl.min = np.array([-2.0, -2.0])
 ctrl.max = np.array([2.0, 2.0])
 ctrl.guess = np.zeros((n, n_u))
 
-states = [qpos, qvel]
-controls = [ctrl]
-
-# ── Dynamics: position kinematics symbolically, velocity via MJX ──────────────
-dynamics: dict = {"qpos": qvel}  # nq == nv
-
-byof: ByofSpec = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-
 # ── Constraints (CTCS on state / control bounds) ─────────────────────────────
 constraints = []
 
@@ -179,13 +171,12 @@ time = ox.Time(
 )
 
 problem = Problem(
-    dynamics=dynamics,
-    states=states,
-    controls=controls,
+    dynamics=dyn,
+    states=dyn.states,
+    controls=dyn.controls,
     time=time,
     constraints=constraints,
     N=n,
-    byof=byof,
     algorithm={
         "lam_prox": 1e-1,
         "lam_cost": 0e0,

{openscvx-2.dev4 → openscvx-2.dev5}/examples/mjx/triple_cartpole_mjx.py

@@ -35,8 +35,7 @@ except ImportError:
     sys.exit(1)
 
 import openscvx as ox
-from openscvx import ByofSpec, Problem
-from openscvx.integrations import mjx_byof
+from openscvx import Problem
 
 L1, L2, L3 = 0.5, 0.4, 0.3  # link lengths (m)
 
@@ -88,35 +87,28 @@ n_u = int(mjx_model.nu)  # 1: cart force
 n = 60  # more nodes → finer resolution near the unstable upright equilibrium
 total_time = 2.5
 
-# ── State / control definitions ───────────────────────────────────────────────
-qpos = ox.State("qpos", shape=(n_q,))
+# ── MJX dynamics adapter ──────────────────────────────────────────────────────
+dyn = ox.MjxDynamics(mjx_model)
+qpos, qvel = dyn.states
+(ctrl,) = dyn.controls
+
 qpos.min = np.array([-100.0, -2 * np.pi, -2 * np.pi, -2 * np.pi])
 qpos.max = np.array([100.0, 2 * np.pi, 2 * np.pi, 2 * np.pi])
 qpos.initial = np.array([0.0, np.pi, 0.0, 0.0])  # all links hanging down
 qpos.final = [ox.Free(0.0), 0.0, 0.0, 0.0]  # all links upright
 
-qvel = ox.State("qvel", shape=(n_v,))
 qvel.min = np.array([-12.0, -12.0, -12.0, -12.0])
 qvel.max = np.array([12.0, 12.0, 12.0, 12.0])
 qvel.initial = np.zeros(n_v)
 qvel.final = [0.0, 0.0, 0.0, 0.0]
 
-ctrl = ox.Control("ctrl", shape=(n_u,))
 ctrl.min = np.array([-3.0])
 ctrl.max = np.array([3.0])
 ctrl.guess = np.zeros((n, n_u))
 
-states = [qpos, qvel]
-controls = [ctrl]
-
-# ── Dynamics: position kinematics symbolically, velocity via MJX ──────────────
-dynamics: dict = {"qpos": qvel}  # nq==nv so this is always valid
-
-byof: ByofSpec = {"dynamics": mjx_byof(mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)}
-
 # ── Constraints (CTCS on state / control bounds) ───────────────────────────────
 constraints = []
-for state in states:
+for state in dyn.states:
     constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
 
 # ── Initial guess: linearly swing θ₁ from π → 0, others stay 0 ───────────────
@@ -132,13 +124,12 @@ time = ox.Time(
 )
 
 problem = Problem(
-    dynamics=dynamics,
-    states=states,
-    controls=controls,
+    dynamics=dyn,
+    states=dyn.states,
+    controls=dyn.controls,
     time=time,
     constraints=constraints,
     N=n,
-    byof=byof,
     algorithm={
         "lam_prox": 1e0,
         "lam_cost": 0e0,

{openscvx-2.dev4 → openscvx-2.dev5}/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-2.dev4 → openscvx-2.dev5}/openscvx/_version.py

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

openscvx-2.dev5/openscvx/integrations/__init__.py

@@ -0,0 +1,66 @@
+"""External-backend dynamics adapters for OpenSCvx.
+
+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 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, ...)
+"""
+
+from typing import Any
+
+from .base import DynamicsAdapter
+from .mjx import MjxDynamics
+
+
+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)
+
+
+def __getattr__(name: str) -> Any:
+    if name == "menagerie":
+        from . import menagerie
+
+        return menagerie
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "DynamicsAdapter",
+    "MjxDynamics",
+    "mjx_dynamics",
+    "menagerie",
+]

openscvx-2.dev5/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