mujoco-mojo 2.3.2__tar.gz → 2.3.3__tar.gz

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/.gitignore

@@ -26,3 +26,5 @@ MUJOCO_LOG.TXT
 # TypeScript build artifacts (node_modules is dev-only; compiled .js files ARE committed)
 node_modules/
 *.js.map
+
+.coverage

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mujoco-mojo
-Version: 2.3.2
+Version: 2.3.3
 Summary: A complete MJCF lifecycle and trial orchestration suite for MuJoCo, powered by Pydantic v2.
 Project-URL: Homepage, https://github.com/Hydrowelder/mujoco-mojo
 Project-URL: Documentation, https://hydrowelder.github.io/mujoco-mojo/

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/docs/user-guides/getting-started.md

@@ -43,14 +43,13 @@ Before writing code, it is helpful to understand the three pillars of a Mojo pro
         def runtime(
             mojo_model: mojo.MojoModel,
             runtime_manager: rt.RuntimeManager,
-            mj_model: mujoco.MjModel,
-            mj_data: mujoco.MjData,
+            state: mojo.MjState,
             *args,
             **kwargs,
         ):
             # Step the physics and apply logic here
-            while mj_data.time < 10.0:
-                runtime_manager.step(mj_model, mj_data)
+            while state.data.time < 10.0:
+                runtime_manager.step(state)
         ```
 
 - **Request-Based Telemetry:** You don't (*have to*) manually log data. Instead, you "Request" that specific geoms, sites, or custom forces be tracked by the `SignalManager`.

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/docs/user-guides/monte_carlo_example.py

@@ -2,7 +2,6 @@
 from pathlib import Path
 from typing import Literal
 
-import mujoco
 import numpy as np
 from pydantic import Field
 
@@ -254,8 +253,7 @@ def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel:
 def runtime(
     mojo_model: mojo.MojoModel,
     runtime_manager: rt.RuntimeManager,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
     *args,
     **kwargs,
 ) -> mojo.MojoModel:
@@ -272,7 +270,7 @@ def runtime(
                 path=Path("runtime-anim.gif"),
                 camera_name=FIXED_CAMERA_NAME,
                 show_loads=True,
-            ).setup(mj_model).register_to_rm(rm)
+            ).setup(state).register_to_rm(rm)
         # --8<-- [end:video]
 
         # Register forces from our handoff data
@@ -293,8 +291,8 @@ def runtime(
 
         # --8<-- [start:stepping]
         # Step until 2.0 seconds
-        while mj_data.time < 2.0:
-            rm.step(mj_model, mj_data)
+        while state.data.time < 2.0:
+            rm.step(state)
         # --8<-- [end:stepping]
 
     return mojo_model

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/docs/user-guides/optimization.md

@@ -54,7 +54,7 @@ The **Objective Function** is the "grade" you give to a simulation. It is a stan
     --8<-- "docs/user-guides/optimization_example.py:score-handle"
     ```
 
-A valid objective function must accept the `MojoModel`, the `Path` to the telemetry outputs requested during runtime, and the MuJoCo `mj_model`/`mj_data` objects.
+A valid objective function must accept the `MojoModel`, the `Path` to the telemetry outputs requested during runtime, and an `MjState` object.
 
 ???+ example "Example: Scoring Performance"
 

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/docs/user-guides/optimization_example.py

@@ -1,7 +1,6 @@
 from pathlib import Path
 from typing import Literal
 
-import mujoco
 import numpy as np
 from pydantic import Field
 
@@ -196,8 +195,7 @@ def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel:
 def runtime(
     mojo_model: mojo.MojoModel,
     runtime_manager: rt.RuntimeManager,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
     *args,
     **kwargs,
 ) -> mojo.MojoModel:
@@ -215,8 +213,8 @@ def runtime(
                 b.request(rm.signal_manager)
             handoff.box1_rot.request(rm.signal_manager)
 
-        while mj_data.time < 2.0:
-            rm.step(mj_model, mj_data)
+        while state.data.time < 2.0:
+            rm.step(state)
 
     return mojo_model
 
@@ -226,8 +224,7 @@ def runtime(
 def objective(
     mojo_model: mojo.MojoModel,
     telemetry: Path,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
 ) -> float:
     # --8<-- [end:score-handle]
     """
@@ -235,12 +232,12 @@ def objective(
     """
     handoff = mojo_model.get_user_data(Handoff)
 
-    w1 = handoff.box1.rt_ang_vel(mj_model, mj_data)
-    w2 = handoff.box2.rt_ang_vel(mj_model, mj_data)
+    w1 = handoff.box1.rt_ang_vel(state)
+    w2 = handoff.box2.rt_ang_vel(state)
     omega_score = float(np.linalg.norm(w1 - w2))
 
-    ke1 = handoff.box1.rt_trans_ke(mj_model, mj_data)
-    ke2 = handoff.box2.rt_trans_ke(mj_model, mj_data)
+    ke1 = handoff.box1.rt_trans_ke(state)
+    ke2 = handoff.box2.rt_trans_ke(state)
     total_ke = ke1 + ke2
     ke_score = 1 / (total_ke + 1e-6)
 

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/examples/boxes_and_springs_monte_carlo.py

@@ -1,7 +1,6 @@
 from pathlib import Path
 from typing import Literal
 
-import mujoco
 import numpy as np
 from pydantic import Field
 
@@ -220,8 +219,7 @@ def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel:
 def runtime(
     mojo_model: mojo.MojoModel,
     runtime_manager: rt.RuntimeManager,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
     *args,
     **kwargs,
 ) -> mojo.MojoModel:
@@ -234,7 +232,7 @@ def runtime(
                 path=Path("runtime-anim.gif"),
                 camera_name=FIXED_CAMERA_NAME,
                 show_loads=True,
-            ).setup(mj_model).register_to_rm(rm)
+            ).setup(state).register_to_rm(rm)
 
         # Register forces from our handoff data
         handoff.add_spring_force("pz", rm)
@@ -249,8 +247,8 @@ def runtime(
                 )
             handoff.box1_rot.request(rm.signal_manager)
         # Step until 2.0 seconds
-        while mj_data.time < 2.0:
-            rm.step(mj_model, mj_data)
+        while state.data.time < 2.0:
+            rm.step(state)
 
     return mojo_model
 

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/examples/boxes_and_springs_optimization.py

@@ -1,7 +1,6 @@
 from pathlib import Path
 from typing import Literal
 
-import mujoco
 import numpy as np
 from pydantic import Field
 
@@ -196,8 +195,7 @@ def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel:
 def runtime(
     mojo_model: mojo.MojoModel,
     runtime_manager: rt.RuntimeManager,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
     *args,
     **kwargs,
 ) -> mojo.MojoModel:
@@ -215,8 +213,8 @@ def runtime(
                 b.request(rm.signal_manager)
             handoff.box1_rot.request(rm.signal_manager)
 
-        while mj_data.time < 2.0:
-            rm.step(mj_model, mj_data)
+        while state.data.time < 2.0:
+            rm.step(state)
 
     return mojo_model
 
@@ -224,20 +222,19 @@ def runtime(
 def objective(
     mojo_model: mojo.MojoModel,
     telemetry: Path,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
 ) -> float:
     """
     Score the trial: Low relative angular velocity (stability) weighted against high translational kinetic energy.
     """
     handoff = mojo_model.get_user_data(Handoff)
 
-    w1 = handoff.box1.rt_ang_vel(mj_model, mj_data)
-    w2 = handoff.box2.rt_ang_vel(mj_model, mj_data)
+    w1 = handoff.box1.rt_ang_vel(state)
+    w2 = handoff.box2.rt_ang_vel(state)
     omega_score = float(np.linalg.norm(w1 - w2))
 
-    ke1 = handoff.box1.rt_trans_ke(mj_model, mj_data)
-    ke2 = handoff.box2.rt_trans_ke(mj_model, mj_data)
+    ke1 = handoff.box1.rt_trans_ke(state)
+    ke2 = handoff.box2.rt_trans_ke(state)
     total_ke = ke1 + ke2
     ke_score = 1 / (total_ke + 1e-6)
 

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/examples/tennis_racket_theorem.py

@@ -5,7 +5,6 @@ This file generates an example model of the [tennis racket theorem](https://en.w
 
 from pathlib import Path
 
-import mujoco
 import numpy as np
 
 import mujoco_mojo as mojo
@@ -138,8 +137,7 @@ def generate(mojo_model: mojo.MojoModel, *args, **kwargs) -> mojo.MojoModel:
 def runtime(
     mojo_model: mojo.MojoModel,
     runtime_manager: rt.RuntimeManager,
-    mj_model: mujoco.MjModel,
-    mj_data: mujoco.MjData,
+    state: mojo.MjState,
     *args,
     **kwargs,
 ) -> mojo.MojoModel:
@@ -148,24 +146,23 @@ def runtime(
     handoff = mojo_model.get_user_data(Handoff)
 
     handoff.racket.set_initial_velocity(
-        mj_model,
-        mj_data,
+        state,
         angular_velocity=np.asarray([0.1, 4.0, 0.0]),
         angle=mojo.Angle.RADIAN,
     )
 
     # Define a custom signal sampler
-    def sample_nutation(mj_model: mujoco.MjModel, mj_data: mujoco.MjData):
+    def sample_nutation(state: mojo.MjState):
         """Adds the [nutation angle](https://en.wikipedia.org/wiki/Nutation) as an output signal."""
         assert runtime_manager.signal_manager
 
         # Get the angular momentum vector and magnitude
-        L = handoff.racket.rt_ang_mom(mj_model, mj_data)
+        L = handoff.racket.rt_ang_mom(state)
         L_mag = np.linalg.norm(L)
 
         if L_mag > 1e-9:
             # Get the handle axis in the world frame
-            R = handoff.racket.rt_xmat(mj_model, mj_data)
+            R = handoff.racket.rt_xmat(state)
             handle_axis = R[:, 0]
 
             # Calculate the angle (theta) between L and the handle axis
@@ -191,8 +188,8 @@ def runtime(
                 b.request(rm.signal_manager)
 
         # Run to observe multiple flips
-        while mj_data.time < 10.0:
-            rm.step(mj_model, mj_data)
+        while state.data.time < 10.0:
+            rm.step(state)
 
     return mojo_model
 

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/src/mujoco_mojo/__init__.py

@@ -198,6 +198,7 @@ from .mjcf import (
     XYAxes,
     ZAxis,
 )
+from .mj_state import MjState
 from .mojo_model import MojoModel, UserData
 from .stochas import (
     NOMINAL_TRIAL_NUM,
@@ -467,6 +468,7 @@ __all__ = [
     "MeshSupersphere",
     "MeshTorus",
     "MeshWedge",
+    "MjState",
     "Model",
     "ModelName",
     "MojoModel",

mujoco_mojo-2.3.3/src/mujoco_mojo/mj_state.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import mujoco
+
+__all__ = ["MjState"]
+
+
+@dataclass
+class MjState:
+    """Pairs MjModel and MjData for passing through runtime methods."""
+
+    model: mujoco.MjModel
+    data: mujoco.MjData

{mujoco_mojo-2.3.2 → mujoco_mojo-2.3.3}/src/mujoco_mojo/mjcf/mujoco.py

@@ -5,6 +5,8 @@ from pathlib import Path
 import mujoco
 from pydantic import Field
 
+from mujoco_mojo.mj_state import MjState
+
 import mujoco_mojo.utils.utils as utils
 from mujoco_mojo.mjcf.defaults import DEFAULT_ANGLE, DEFAULT_EULERSEQ
 from mujoco_mojo.mjcf.extension import Extension
@@ -242,16 +244,14 @@ class Mujoco(XMLModel):
             )
         )
 
-    def prep_for_sim(
-        self, save_path: Path | None = None
-    ) -> tuple[mujoco.MjModel, mujoco.MjData]:
-        """Creates a MuJoCo MjModel and MjData from the Mujoco instance."""
+    def prep_for_sim(self, save_path: Path | None = None) -> MjState:
+        """Creates an MjState (MjModel + MjData) from the Mujoco instance."""
         if save_path:
             self.write_xml(save_path)
-            mj_model = mujoco.MjModel.from_xml_path(str(save_path))
+            model = mujoco.MjModel.from_xml_path(str(save_path))
         else:
-            mj_model = self.to_mj_model()
+            model = self.to_mj_model()
 
-        mj_data = mujoco.MjData(mj_model)
-        mujoco.mj_forward(mj_model, mj_data)
-        return mj_model, mj_data
+        data = mujoco.MjData(model)
+        mujoco.mj_forward(model, data)
+        return MjState(model, data)