openscvx 0.5.2.dev2__py3-none-any.whl → 0.5.2.dev12__py3-none-any.whl

openscvx/__init__.py

@@ -12,6 +12,7 @@ import openscvx.symbolic.expr.spatial as spatial
 import openscvx.symbolic.expr.stl as stl
 import openscvx.symbolic.expr.stljax as stljax
 from openscvx.algorithms import (
+    AdaptiveProximalWeight,
     AugmentedLagrangian,
     ConstantProximalWeight,
     PenalizedTrustRegion,
@@ -25,9 +26,13 @@ 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
+from openscvx.solvers import CVXPyPTRSolver, PTRSolver
+
+# QPAXPTRSolver is exposed lazily via __getattr__ below to keep `import qpax`
+# off the hot import path for users who don't install the optional extra.
 from openscvx.symbolic.expr import (
     CTCS,
     Abs,
@@ -89,6 +94,20 @@ from openscvx.utils.cache import clear_cache, get_cache_dir, get_cache_size
 
 load_results = OptimizationResults.load
 
+
+def __getattr__(name: str):
+    """Lazy export for backends that depend on optional packages."""
+    if name == "QPAXPTRSolver":
+        from openscvx.solvers.qpax_ptr_solver import QPAXPTRSolver
+
+        return QPAXPTRSolver
+    if name == "MoreauPTRSolver":
+        from openscvx.solvers.moreau_ptr_solver import MoreauPTRSolver
+
+        return MoreauPTRSolver
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     # Main Trajectory Optimization Entrypoint
     "Problem",
@@ -176,6 +195,9 @@ __all__ = [
     "lie",
     # Expert mode types
     "ByofSpec",
+    # External-backend dynamics adapters
+    "DynamicsAdapter",
+    "MjxDynamics",
     # Discretization
     "DiscretizeLinearizeVectorize",
     "LinearizeDiscretize",
@@ -183,9 +205,13 @@ __all__ = [
     "VectorizeDiscretizeLinearize",
     # Convex Solver
     "PTRSolver",
+    "CVXPyPTRSolver",
+    "QPAXPTRSolver",
+    "MoreauPTRSolver",
     # Algorithm & Autotuning
     "PenalizedTrustRegion",
     "AugmentedLagrangian",
+    "AdaptiveProximalWeight",
     "ConstantProximalWeight",
     "RampProximalWeight",
 ]

openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.2.dev2'
-__version_tuple__ = version_tuple = (0, 5, 2, 'dev2')
+__version__ = version = '0.5.2.dev12'
+__version_tuple__ = version_tuple = (0, 5, 2, 'dev12')
 
 __commit_id__ = commit_id = None

openscvx/algorithms/__init__.py

@@ -82,12 +82,19 @@ from typing import Annotated, Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 
-from .augmented_lagrangian import AugmentedLagrangian, AugmentedLagrangianSpec
+from .autotuner import (
+    AdaptiveProximalWeight,
+    AdaptiveProximalWeightSpec,
+    AugmentedLagrangian,
+    AugmentedLagrangianSpec,
+    ConstantProximalWeight,
+    ConstantProximalWeightSpec,
+    RampProximalWeight,
+    RampProximalWeightSpec,
+)
 from .base import Algorithm, AlgorithmState, AutotuningBase, DiscretizationResult
-from .constant_proximal_weight import ConstantProximalWeight, ConstantProximalWeightSpec
 from .optimization_results import OptimizationResults
-from .penalized_trust_region import PenalizedTrustRegion
-from .ramp_proximal_weight import RampProximalWeight, RampProximalWeightSpec
+from .scvx import PenalizedTrustRegion
 from .weights import Weights
 
 # ---------------------------------------------------------------------------
@@ -97,6 +104,7 @@ from .weights import Weights
 AutotunerConfig = Annotated[
     Union[
         AugmentedLagrangianSpec,
+        AdaptiveProximalWeightSpec,
         RampProximalWeightSpec,
         ConstantProximalWeightSpec,
     ],
@@ -173,6 +181,7 @@ __all__ = [
     "PenalizedTrustRegion",
     "AutotuningBase",
     "AugmentedLagrangian",
+    "AdaptiveProximalWeight",
     "ConstantProximalWeight",
     "RampProximalWeight",
     # Config models

openscvx/algorithms/autotuner/__init__.py

@@ -0,0 +1,17 @@
+"""SCP weight autotuning strategies."""
+
+from .adaptive_proximal_weight import AdaptiveProximalWeight, AdaptiveProximalWeightSpec
+from .augmented_lagrangian import AugmentedLagrangian, AugmentedLagrangianSpec
+from .constant_proximal_weight import ConstantProximalWeight, ConstantProximalWeightSpec
+from .ramp_proximal_weight import RampProximalWeight, RampProximalWeightSpec
+
+__all__ = [
+    "AdaptiveProximalWeight",
+    "AdaptiveProximalWeightSpec",
+    "AugmentedLagrangian",
+    "AugmentedLagrangianSpec",
+    "ConstantProximalWeight",
+    "ConstantProximalWeightSpec",
+    "RampProximalWeight",
+    "RampProximalWeightSpec",
+]

openscvx/algorithms/autotuner/adaptive_proximal_weight.py

@@ -0,0 +1,190 @@
+"""Autotuning functions for SCP (Successive Convex Programming) parameters."""
+
+from copy import deepcopy
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict
+
+from openscvx.config import Config
+
+from ..base import AutotuningBase
+from .augmented_lagrangian import AugmentedLagrangian
+
+if TYPE_CHECKING:
+    from openscvx.lowered import LoweredJaxConstraints
+
+    from ..base import AlgorithmState, CandidateIterate
+    from ..weights import Weights
+
+
+class AdaptiveProximalWeight(AutotuningBase):
+    """PTR-style proximal adaptation with fixed virtual penalty weights.
+
+    Same acceptance-ratio logic as :class:`AugmentedLagrangian` for ``lam_prox``,
+    but ``lam_vc`` and ``lam_vb_*`` are held constant at their current state values.
+    """
+
+    COLUMNS = AugmentedLagrangian.COLUMNS
+
+    def __init__(
+        self,
+        gamma_1: float = 2.0,
+        gamma_2: float = 0.5,
+        eta_0: float = 1e-2,
+        eta_1: float = 1e-1,
+        eta_2: float = 0.8,
+        lam_prox_min: float = 1e-3,
+        lam_prox_max: float = 1e4,
+        lam_cost_drop: int = -1,
+        lam_cost_relax: float = 1.0,
+    ):
+        self.gamma_1 = gamma_1
+        self.gamma_2 = gamma_2
+        self.eta_0 = eta_0
+        self.eta_1 = eta_1
+        self.eta_2 = eta_2
+        self.lam_prox_min = lam_prox_min
+        self.lam_prox_max = lam_prox_max
+        self.lam_cost_drop = lam_cost_drop
+        self.lam_cost_relax = lam_cost_relax
+
+    @staticmethod
+    def _copy_virtual_weights(
+        candidate: "CandidateIterate",
+        state: "AlgorithmState",
+    ) -> None:
+        candidate.lam_vc = state.lam_vc
+        candidate.lam_vb_nodal = state.lam_vb_nodal
+        candidate.lam_vb_cross = state.lam_vb_cross
+
+    def update_weights(
+        self,
+        state: "AlgorithmState",
+        candidate: "CandidateIterate",
+        nodal_constraints: "LoweredJaxConstraints",
+        settings: Config,
+        params: dict,
+        weights: "Weights",
+    ) -> str:
+        """Update SCP proximal weight based on acceptance ratio; keep VC/VB fixed."""
+        candidate_x_prop = (
+            candidate.x_prop_plus[1:] if candidate.x_prop_plus is not None else candidate.x_prop
+        )
+        (
+            nonlinear_cost,
+            nonlinear_penalty,
+            nodal_penalty,
+        ) = self.calculate_nonlinear_penalty(
+            candidate_x_prop,
+            candidate.x,
+            candidate.u,
+            state.lam_vc,
+            state.lam_vb_nodal,
+            state.lam_vb_cross,
+            state.lam_cost,
+            nodal_constraints,
+            params,
+            settings,
+        )
+
+        candidate.J_nonlin = nonlinear_cost + nonlinear_penalty + nodal_penalty
+
+        if state.k > self.lam_cost_drop:
+            candidate.lam_cost = state.lam_cost * self.lam_cost_relax
+        else:
+            candidate.lam_cost = weights.lam_cost
+
+        lam_prox_k = deepcopy(state.lam_prox)
+
+        if state.k > 1:
+            state_x_prop_plus = state.x_prop_plus()
+            state_x_prop = (
+                state_x_prop_plus[1:] if state_x_prop_plus is not None else state.x_prop()
+            )
+            (
+                prev_nonlinear_cost,
+                prev_nonlinear_penalty,
+                prev_nodal_penalty,
+            ) = self.calculate_nonlinear_penalty(
+                state_x_prop,
+                state.x,
+                state.u,
+                state.lam_vc,
+                state.lam_vb_nodal,
+                state.lam_vb_cross,
+                state.lam_cost,
+                nodal_constraints,
+                params,
+                settings,
+            )
+
+            J_nonlin_prev = prev_nonlinear_cost + prev_nonlinear_penalty + prev_nodal_penalty
+
+            actual_reduction = J_nonlin_prev - candidate.J_nonlin
+            predicted_reduction = J_nonlin_prev - candidate.J_lin
+
+            if predicted_reduction == 0:
+                raise ValueError("Predicted reduction is 0.")
+
+            rho = actual_reduction / predicted_reduction
+
+            state.pred_reduction_history.append(predicted_reduction)
+            state.actual_reduction_history.append(actual_reduction)
+            state.acceptance_ratio_history.append(rho)
+
+            if rho < self.eta_0:
+                lam_prox_k1 = np.minimum(self.lam_prox_max, self.gamma_1 * lam_prox_k)
+                candidate.lam_prox = lam_prox_k1
+                state.reject_solution(candidate)
+                adaptive_state = "Reject Higher"
+            elif rho >= self.eta_0 and rho < self.eta_1:
+                lam_prox_k1 = np.minimum(self.lam_prox_max, self.gamma_1 * lam_prox_k)
+                candidate.lam_prox = lam_prox_k1
+                self._copy_virtual_weights(candidate, state)
+                state.accept_solution(candidate)
+                adaptive_state = "Accept Higher"
+            elif rho >= self.eta_1 and rho < self.eta_2:
+                candidate.lam_prox = lam_prox_k
+                self._copy_virtual_weights(candidate, state)
+                state.accept_solution(candidate)
+                adaptive_state = "Accept Constant"
+            else:
+                lam_prox_k1 = np.maximum(self.lam_prox_min, self.gamma_2 * lam_prox_k)
+                candidate.lam_prox = lam_prox_k1
+                self._copy_virtual_weights(candidate, state)
+                state.accept_solution(candidate)
+                adaptive_state = "Accept Lower"
+
+        else:
+            candidate.lam_prox = lam_prox_k
+            self._copy_virtual_weights(candidate, state)
+            state.accept_solution(candidate)
+            adaptive_state = "Initial"
+
+        return adaptive_state
+
+
+# =============================================================================
+# Pydantic spec for dict / YAML validation
+# =============================================================================
+
+
+class AdaptiveProximalWeightSpec(BaseModel):
+    """Validates AdaptiveProximalWeight configuration from dict/YAML input."""
+
+    type: Literal["AdaptiveProximalWeight"] = "AdaptiveProximalWeight"
+    gamma_1: float = 2.0
+    gamma_2: float = 0.5
+    eta_0: float = 1e-2
+    eta_1: float = 1e-1
+    eta_2: float = 0.8
+    lam_prox_min: float = 1e-3
+    lam_prox_max: float = 1e4
+    lam_cost_drop: int = -1
+    lam_cost_relax: float = 1.0
+
+    model_config = ConfigDict(extra="forbid")
+
+    def build(self) -> AdaptiveProximalWeight:
+        return AdaptiveProximalWeight(**self.model_dump(exclude={"type"}, exclude_unset=True))

openscvx/algorithms/{augmented_lagrangian.py → autotuner/augmented_lagrangian.py}

@@ -14,13 +14,13 @@ from openscvx.utils.printing import (
     color_adaptive_state,
 )
 
-from .base import AutotuningBase
+from ..base import AutotuningBase
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
 
-    from .base import AlgorithmState, CandidateIterate
-    from .weights import Weights
+    from ..base import AlgorithmState, CandidateIterate
+    from ..weights import Weights
 
 
 class AugmentedLagrangian(AutotuningBase):
@@ -86,8 +86,12 @@ class AugmentedLagrangian(AutotuningBase):
             eta_2: Threshold above which solution is accepted with lower weight.
                 Defaults to 0.8.
             ep: Threshold for virtual control weight update (nu > ep vs nu <= ep).
-                Defaults to 0.5.
-            eta_lambda: Step size for virtual control weight update. Defaults to 1e0.
+                Must lie in (0, 1). Defaults to 0.99; when tuning, try 1e-1 if needed.
+                Typically tuned together with ``eta_lambda`` (often the first
+                parameters adjusted).
+            eta_lambda: Step size for virtual control weight update. Defaults to 1e1.
+                Typically tuned together with ``ep`` (often the first parameters
+                adjusted).
             lam_vc_max: Maximum virtual control penalty weight. Defaults to 1e5.
             lam_prox_min: Minimum trust region (proximal) weight. Defaults to 1e-3.
             lam_prox_max: Maximum trust region (proximal) weight. Defaults to 2e5.

openscvx/algorithms/{constant_proximal_weight.py → autotuner/constant_proximal_weight.py}

@@ -6,13 +6,13 @@ from pydantic import BaseModel, ConfigDict
 
 from openscvx.config import Config
 
-from .base import AutotuningBase
+from ..base import AutotuningBase
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
 
-    from .base import AlgorithmState, CandidateIterate
-    from .weights import Weights
+    from ..base import AlgorithmState, CandidateIterate
+    from ..weights import Weights
 
 
 class ConstantProximalWeight(AutotuningBase):

openscvx/algorithms/{ramp_proximal_weight.py → autotuner/ramp_proximal_weight.py}

@@ -7,13 +7,13 @@ from pydantic import BaseModel, ConfigDict
 
 from openscvx.config import Config
 
-from .base import AutotuningBase
+from ..base import AutotuningBase
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
 
-    from .base import AlgorithmState, CandidateIterate
-    from .weights import Weights
+    from ..base import AlgorithmState, CandidateIterate
+    from ..weights import Weights
 
 
 class RampProximalWeight(AutotuningBase):

openscvx/algorithms/scvx/__init__.py

@@ -0,0 +1,5 @@
+"""Successive convexification algorithm implementations."""
+
+from .penalized_trust_region import PenalizedTrustRegion
+
+__all__ = ["PenalizedTrustRegion"]

openscvx/algorithms/{penalized_trust_region.py → scvx/penalized_trust_region.py}

@@ -6,8 +6,9 @@ optimization problems through iterative convex approximation.
 
 import time
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, List, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Tuple, Union
 
+import jax
 import numpy as np
 import numpy.linalg as la
 
@@ -21,11 +22,11 @@ from openscvx.utils.printing import (
     color_prob_stat,
 )
 
-from .augmented_lagrangian import AugmentedLagrangian
-from .base import Algorithm, AlgorithmState, CandidateIterate
-from .constant_proximal_weight import ConstantProximalWeight
-from .ramp_proximal_weight import RampProximalWeight
-from .weights import Weights
+from ..autotuner.augmented_lagrangian import AugmentedLagrangian
+from ..autotuner.constant_proximal_weight import ConstantProximalWeight
+from ..autotuner.ramp_proximal_weight import RampProximalWeight
+from ..base import Algorithm, AlgorithmState, CandidateIterate
+from ..weights import Weights
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
@@ -33,7 +34,7 @@ if TYPE_CHECKING:
     from openscvx.symbolic.expr.control import Control
     from openscvx.symbolic.expr.state import State
 
-    from .base import AutotuningBase
+    from ..base import AutotuningBase
 
 warnings.filterwarnings("ignore")
 
@@ -188,6 +189,11 @@ class PenalizedTrustRegion(Algorithm):
             return solver.call(*args)
         return solver(*args)
 
+    @staticmethod
+    def _block_until_ready_outputs(outputs: Tuple[object, ...]) -> None:
+        """Finish any pending XLA work from discretization exports (warm-up helper)."""
+        jax.block_until_ready(outputs)
+
     def _recover_prior_node_from_initial(
         self,
         settings: Config,
@@ -286,7 +292,8 @@ class PenalizedTrustRegion(Algorithm):
         """Initialize PTR algorithm.
 
         Stores compiled infrastructure and performs a warm-start solve to
-        initialize DPP and JAX jacobians.
+        initialize DPP and JAX jacobians. Also runs post-subproblem discretization
+        on the throwaway CVX solution so XLA/export caches match the first ``step()``.
 
         Args:
             solver: Convex subproblem solver (e.g., CVXPySolver)
@@ -333,7 +340,29 @@ class PenalizedTrustRegion(Algorithm):
                 params,
             )
             init_state.add_impulsive_discretization(W_multi_shoot.__array__())
-        _ = self._subproblem(params, init_state, settings)
+        (x_sol, u_sol, *_) = self._subproblem(params, init_state, settings)
+
+        # Prime the same exported discretization calls used after every subproblem in
+        # step() (candidate trajectory). initialize() previously only discretized the
+        # initial guess, so the first step() in solve() could still hit an XLA cache_miss
+        # on post-CVX (x_sol, u_sol). Running that path here moves compilation into init.
+        cont_out = self._invoke_solver(
+            self._discretization_solver, x_sol, u_sol.astype(float), params
+        )
+        x_prop_c = cont_out[3]
+        u_candidate = u_sol.astype(float)
+        x0_prior_c = self._recover_prior_node_from_initial(settings, x_sol[0])
+        x_nodes_prior_c = np.vstack((x0_prior_c, np.asarray(x_prop_c)))
+        if self._discretization_solver_impulsive is not None:
+            imp_out = self._invoke_solver(
+                self._discretization_solver_impulsive,
+                x_nodes_prior_c,
+                u_candidate,
+                params,
+            )
+            self._block_until_ready_outputs(cont_out + imp_out)
+        else:
+            self._block_until_ready_outputs(cont_out)
 
     def step(
         self,

openscvx/discretization/discretize_linearize.py

@@ -188,7 +188,12 @@ class VectorizeDiscretizeLinearize(Discretizer):
                 in_axes=-1,
                 out_axes=-1,
             )
-            jacobians = pushforward(standard_basis)
+            # Cast to match primal's dtype.  standard_basis is built eagerly
+            # from jnp.eye (float32 when x64 is disabled), but jax.export
+            # traces with the literal dtype of the dummy arrays (float64 from
+            # np.ones), so primal may be float64 while the closure-captured
+            # standard_basis is float32.  jax.jvp requires identical dtypes.
+            jacobians = pushforward(standard_basis.astype(primal.dtype))
 
             A_d = jacobians[:, :, :, i0:i1]
             B_d = jacobians[:, :, :, i1:i2]

openscvx/discretization/sparse_utils/sparse_jacobian.py

@@ -125,7 +125,12 @@ def _sparse_jacobian_fn(
             _, jvp_out = jax.jvp(f_of_target, (primals[argnums],), (tangent,))
             return jvp_out
 
-        compressed = jax.vmap(single_jvp)(seeds)  # (n_colors, out_dim)
+        # Cast seeds to match the primal dtype.  seeds is built eagerly as
+        # float32 (x64 disabled), but jax.export traces with the literal
+        # dtype of the dummy arrays (float64 from np.ones), so primals may
+        # be float64.  jax.jvp requires matching dtypes.
+        typed_seeds = seeds.astype(primals[argnums].dtype)
+        compressed = jax.vmap(single_jvp)(typed_seeds)  # (n_colors, out_dim)
 
         values = compressed[scatter_color, scatter_row]
         jac = jnp.zeros((out_dim, in_dim))

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",
+]