mujoco-mojo 2.3.4__tar.gz → 2.3.6__tar.gz

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/.pre-commit-config.yaml

@@ -46,6 +46,12 @@ repos:
     rev: v1.1.408
     hooks:
       - id: pyright
+      - id: pyright
+        name: Pyright (docs examples)
+        args: [docs/user-guides]
+        pass_filenames: false
+        always_run: true
+        stages: [pre-commit]
   - repo: local
     hooks:
       - id: act-test

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mujoco-mojo
-Version: 2.3.4
+Version: 2.3.6
 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/
@@ -45,6 +45,7 @@ Requires-Dist: pybind11-stubgen>=2.5.5
 Requires-Dist: pydantic-settings>=2.14.1
 Requires-Dist: pydantic>=2.12.5
 Requires-Dist: python-fcl>=0.7.0.11
+Requires-Dist: pytransform3d>=3.15.0
 Requires-Dist: rich>=14.3.3
 Requires-Dist: rtree>=1.4.1
 Requires-Dist: scipy-stubs[scipy]>=1.17.0.2

mujoco_mojo-2.3.6/docs/user-guides/pose-context.md

@@ -0,0 +1,57 @@
+# Pose Context
+
+!!! abstract
+    MJCF poses are always relative to a parent body. **Pose Context** lets you express any element's pose relative to a frame on a different branch of the kinematic tree, without walking the chain manually.
+
+---
+
+## The Problem
+
+When building a kinematic tree, every body and site stores its pose relative to its immediate parent. If you want to place a site on `gripper` at the same world location as `target_site` on `base` (two separate branches) there is no direct MJCF mechanism for it. You would have to compose the full transform chain by hand.
+
+```python
+--8<-- "docs/user-guides/pose_context_example.py:build-tree"
+```
+
+---
+
+## One-Shot Resolution
+
+For a single pose, `Mujoco.local_pose(frame, relative_to)` resolves the pose in one call. The result is a concrete `PoseQuat` you can assign directly to any MJCF element.
+
+```python
+--8<-- "docs/user-guides/pose_context_example.py:one-shot"
+```
+
+`frame` and `relative_to` accept any object that carries a `pose` attribute: `Body`, any site type, `Frame`, `Camera`, `Light`, and so on.
+
+???+ warning "Warning: Match relative_to to the element's parent"
+    `relative_to` must be the direct parent body of the element that will receive the resolved pose. If you assign the result to a site on `arm` but pass `relative_to=gripper`, the pose will be mathematically correct in the wrong coordinate frame and the element will appear in the wrong location.
+
+???+ warning "Warning: Tree must be complete"
+    Both `frame` and `relative_to` must be part of `mojo_model.mjcf.worldbody` at the time of the call. This is the same requirement MuJoCo has for compiling the model - anything not attached to the world will not appear in `MjModel`.
+
+---
+
+## Deferred Resolution with `PoseRef`
+
+`PoseRef` stores the `frame` and `relative_to` pair and resolves lazily by passing `mojo_model.mjcf` to `.to_quat()`. This is useful when you want to define references early and resolve them later, or apply the same reference in multiple places.
+
+```python
+--8<-- "docs/user-guides/pose_context_example.py:batch"
+```
+
+---
+
+## Using `Frame` as a Reference
+
+MJCF `Frame` elements are coordinate transforms that disappear at compile time, accumulating into their children. They are fully supported as `frame` or `relative_to` arguments.
+
+```python
+--8<-- "docs/user-guides/pose_context_example.py:frame-ref"
+```
+
+---
+
+!!! success
+    You can now place elements precisely anywhere in the kinematic tree without manually composing transform chains. Use `local_pose` for quick one-off resolutions and `PoseRef` when you want to define the reference early and resolve it later in your generate function.

mujoco_mojo-2.3.6/docs/user-guides/pose_context_example.py

@@ -0,0 +1,65 @@
+import numpy as np
+
+import mujoco_mojo as mojo
+
+# --8<-- [start:build-tree]
+# `base` and `gripper` are separate bodies, both direct children of worldbody.
+# Neither is a parent of the other (they are on different branches of the tree).
+base = mojo.Body(
+    name=mojo.BodyName("base"), pose=mojo.PoseQuat(pos=np.array([1.0, 0.0, 0.0]))
+)
+
+# `target_site` is on `base`, so its pose ([0, 0, 0.5]) is local to `base`.
+# Its world position is therefore [1.0, 0.0, 0.5].
+target_site = mojo.SiteSphere(
+    name=mojo.SiteName("target"), pose=mojo.PoseQuat(pos=np.array([0.0, 0.0, 0.5]))
+)
+base.sites.append(target_site)
+
+# `gripper` is on a completely different branch. We want to place a site on
+# it at the same world position as `target_site`, but MJCF only accepts poses
+# relative to the direct parent body. There is no built-in way to say
+# "put this site wherever target_site is" across branches.
+gripper = mojo.Body(
+    name=mojo.BodyName("gripper"), pose=mojo.PoseQuat(pos=np.array([-1.0, 0.0, 0.0]))
+)
+
+mojo_model = mojo.MojoModel()
+mojo_model.mjcf.worldbody = mojo.WorldBody()
+mojo_model.mjcf.worldbody.bodies.extend([base, gripper])
+# --8<-- [end:build-tree]
+
+
+# --8<-- [start:one-shot]
+# local_pose takes the objects themselves (not their pose values) because it
+# needs to walk each object's parent chain up to worldbody to compose the
+# full world-space transform. A bare PoseQuat contains no tree context.
+# relative_to must match the parent body of whatever element receives this pose.
+local_pose = mojo_model.mjcf.local_pose(frame=target_site, relative_to=gripper)
+
+mirror = mojo.SiteSphere(name=mojo.SiteName("mirror"), pose=local_pose)
+gripper.sites.append(mirror)
+# --8<-- [end:one-shot]
+
+
+# --8<-- [start:batch]
+# PoseRef stores the reference and resolves it on demand via the mjcf model
+contact_site = mojo.SiteSphere(
+    name=mojo.SiteName("contact"),
+    pose=mojo.PoseRef(frame=target_site, relative_to=gripper).to_quat(mojo_model.mjcf),
+)
+origin_in_gripper = mojo.PoseRef(frame=base, relative_to=gripper).to_quat(
+    mojo_model.mjcf
+)
+# --8<-- [end:batch]
+
+
+# --8<-- [start:frame-ref]
+# Frame elements are also valid references
+alignment_frame = mojo.Frame(pose=mojo.PoseQuat(pos=np.array([0.0, 0.5, 0.0])))
+base.frames.append(alignment_frame)
+
+frame_in_gripper = mojo.PoseRef(frame=alignment_frame, relative_to=gripper).to_quat(
+    mojo_model.mjcf
+)
+# --8<-- [end:frame-ref]

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/mkdocs.yml

@@ -14,20 +14,22 @@ watch:
 nav:
   - Home: index.md
   - User Guide:
-      # - user-guide.md
-      - Getting Started: user-guides/getting-started.md
-      - Initializing a Project: user-guides/init.md
-      - Generate Script: user-guides/generate-script.md
-      - Runtime Script: user-guides/runtime-script.md
-      - Running Jobs:
-          - user-guides/running-jobs.md
-          - Example Report: user-guides/MOJO_RUNTIME_REPORT.md
-      - Dojo Dashboard:
-          - user-guides/dojo.md
-          - Monitor: user-guides/monitor.md
-          - Mosaic: user-guides/mosaic.md
-      - Optimization: user-guides/optimization.md
-      - Reloaded: user-guides/reloaded.md
+      - Workflow:
+          - Getting Started: user-guides/getting-started.md
+          - Initializing a Project: user-guides/init.md
+          - Generate Script: user-guides/generate-script.md
+          - Runtime Script: user-guides/runtime-script.md
+          - Running Jobs:
+              - user-guides/running-jobs.md
+              - Example Report: user-guides/MOJO_RUNTIME_REPORT.md
+      - Features:
+          - Dojo Dashboard:
+              - user-guides/dojo.md
+              - Monitor: user-guides/monitor.md
+              - Mosaic: user-guides/mosaic.md
+          - Optimization: user-guides/optimization.md
+          - Reloaded: user-guides/reloaded.md
+          - Pose Context: user-guides/pose-context.md
   - Code Reference:
       - mujoco_mojo:
           - MojoModel: reference/mujoco_mojo/mojo_model.md

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
     "filelock>=3.29.0",
     "pydantic-settings>=2.14.1",
     "tomli-w>=1.2.0",
+    "pytransform3d>=3.15.0",
 ]
 classifiers = [
     "Development Status :: 4 - Beta",

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/src/mujoco_mojo/__init__.py

@@ -31,6 +31,7 @@ compatibility patch that ships alongside the stubs:
 from mujoco_mojo.__about__ import __version__  # noqa: F401
 
 from . import utils
+from .mj_state import MjState
 from .mjcf import (
     Actuator,
     ActuatorAdhesion,
@@ -93,6 +94,7 @@ from .mjcf import (
     GeomPlane,
     GeomSDF,
     GeomSphere,
+    HasPose,
     HField,
     Inertial,
     Joint,
@@ -117,8 +119,10 @@ from .mjcf import (
     Plugin,
     Pos,
     PoseAxisAngle,
+    PoseContext,
     PoseEuler,
     PoseQuat,
+    PoseRef,
     PoseXYAxes,
     PoseZAxis,
     Quat,
@@ -198,7 +202,6 @@ from .mjcf import (
     XYAxes,
     ZAxis,
 )
-from .mj_state import MjState
 from .mojo_model import MojoModel, UserData
 from .stochas import (
     NOMINAL_TRIAL_NUM,
@@ -426,6 +429,7 @@ __all__ = [
     "GridLayoutStr",
     "HField",
     "HFieldName",
+    "HasPose",
     "Inertia",
     "InertiaFromGeom",
     "InertiaGroupRange",
@@ -485,8 +489,10 @@ __all__ = [
     "PoissonDistribution",
     "Pos",
     "PoseAxisAngle",
+    "PoseContext",
     "PoseEuler",
     "PoseQuat",
+    "PoseRef",
     "PoseXYAxes",
     "PoseZAxis",
     "ProximityType",

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/src/mujoco_mojo/mjcf/__init__.py

@@ -165,6 +165,7 @@ from .orientation import (
 )
 from .plugin import Plugin
 from .pose import AnyPose, PoseAxisAngle, PoseEuler, PoseQuat, PoseXYAxes, PoseZAxis
+from .pose_context import HasPose, PoseContext, PoseRef
 from .position import Pos
 
 __all__ = [
@@ -230,6 +231,7 @@ __all__ = [
     "GeomSDF",
     "GeomSphere",
     "HField",
+    "HasPose",
     "Inertial",
     "Joint",
     "Key",
@@ -253,8 +255,10 @@ __all__ = [
     "Plugin",
     "Pos",
     "PoseAxisAngle",
+    "PoseContext",
     "PoseEuler",
     "PoseQuat",
+    "PoseRef",
     "PoseXYAxes",
     "PoseZAxis",
     "Quat",

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/src/mujoco_mojo/mjcf/mujoco.py

@@ -5,9 +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.mj_state import MjState
 from mujoco_mojo.mjcf.defaults import DEFAULT_ANGLE, DEFAULT_EULERSEQ
 from mujoco_mojo.mjcf.extension import Extension
 from mujoco_mojo.mjcf.mujoco_attr.actuator import Actuator
@@ -24,6 +23,8 @@ from mujoco_mojo.mjcf.mujoco_attr.size import Size
 from mujoco_mojo.mjcf.mujoco_attr.statistic import Statistic
 from mujoco_mojo.mjcf.mujoco_attr.tendon import Tendon
 from mujoco_mojo.mjcf.mujoco_attr.visual import Visual
+from mujoco_mojo.mjcf.pose import PoseQuat
+from mujoco_mojo.mjcf.pose_context import HasPose, PoseContext
 from mujoco_mojo.mjcf.xml_model import XMLModel
 from mujoco_mojo.typing import Angle, EulerSeq, ModelName
 from mujoco_mojo.utils.log import get_logger
@@ -244,6 +245,30 @@ class Mujoco(XMLModel):
             )
         )
 
+    @property
+    def pose_context(self) -> PoseContext:
+        """
+        Builds a PoseContext from the current worldbody tree.
+
+        Raises ValueError if worldbody is not set.
+        """
+        if self.worldbody is None:
+            raise ValueError("Cannot build pose graph: worldbody is not set.")
+        return PoseContext(self.worldbody)
+
+    def local_pose(
+        self,
+        frame: HasPose,
+        relative_to: HasPose,
+    ) -> PoseQuat:
+        """
+        Returns the pose of `frame` expressed in `relative_to`'s coordinate system.
+
+        Always walks the live worldbody tree, so the result reflects the current
+        state of the model regardless of any modifications since the last call.
+        """
+        return self.pose_context.local_pose(frame, relative_to)
+
     def prep_for_sim(self, save_path: Path | None = None) -> MjState:
         """Creates an MjState (MjModel + MjData) from the Mujoco instance."""
         if save_path:

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/src/mujoco_mojo/mjcf/mujoco_attr/body_attr/geom.py

@@ -7,11 +7,11 @@ import numpy as np
 import trimesh
 from pydantic import ConfigDict, Field
 
+from mujoco_mojo.mj_state import MjState
 from mujoco_mojo.mjcf.defaults import SOLIMP_DEFAULT, SOLREF_DEFAULT
 from mujoco_mojo.mjcf.plugin import Plugin
 from mujoco_mojo.mjcf.pose import AnyPose, PoseQuat
 from mujoco_mojo.mjcf.xml_model import XMLModel
-from mujoco_mojo.mj_state import MjState
 from mujoco_mojo.typing import (
     FluidShape,
     GeomName,
@@ -262,7 +262,7 @@ class GeomBase(XMLModel):
         signal_manager: SignalManager,
         attrs: list[
             Literal["xpos", "xmat", "xvelp", "xvelr", "xaccp", "xaccr", "quat"]
-        ] = ["xpos", "quat"],
+        ] = ["xpos", "xvelp", "xvelr", "xaccp", "xaccr", "quat"],
     ):
         """
         Registers specific geom attributes for logging.

mujoco_mojo-2.3.6/src/mujoco_mojo/mjcf/mujoco_attr/sensor_attr/tactile.py

@@ -0,0 +1,18 @@
+from mujoco_mojo.mjcf.mujoco_attr.sensor_attr.base import SensorBase
+from mujoco_mojo.typing import GeomName, MeshName
+
+__all__ = ["SensorTactile"]
+
+
+class SensorTactile(SensorBase):
+    """The tactile sensor returns the maximum penetration depth and the sliding velocities in the tangent frame at given points between the geom associated with the sensor and the SDF geoms in contact with it. The sensor is associated with a geom and a mesh. It is activated by the contact between its associated geom and other geoms. The vertices of the mesh, when positioned in the geom frame, are the points at which sensor values are computed, so the dimension of the output is 3 times the number of vertices in the mesh. The mesh must have 3 normal vectors per vertex, which are used to compute the tangent frame. If the penetration depth is positive (no contact), then all values are 0 for the corresponding vertex. Only contacts with geoms of type SDF contribute to the sensor output. The sensor can be visualized by enabling the visualization of contact points."""
+
+    tag = "tactile"
+
+    attributes = ("name", "user", "geom", "mesh")
+
+    geom: GeomName
+    """Name of the geom to associate the tactile sensor with."""
+
+    mesh: MeshName
+    """Name of the mesh to associate the tactile sensor with. The mesh will be created by the sensor."""

mujoco_mojo-2.3.6/src/mujoco_mojo/mjcf/pose_context.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+import numpy as np
+from pytransform3d.transform_manager import TransformManager
+
+from mujoco_mojo.mjcf.orientation import Quat
+from mujoco_mojo.mjcf.pose import AnyPose, PoseQuat
+from mujoco_mojo.utils.log import get_logger
+
+if TYPE_CHECKING:
+    from mujoco_mojo.mjcf.meta.frame import Frame
+    from mujoco_mojo.mjcf.mujoco import Mujoco
+    from mujoco_mojo.mjcf.mujoco_attr.body import Body, WorldBody
+
+__all__ = ["HasPose", "PoseContext", "PoseRef"]
+
+logger = get_logger(__name__)
+
+_WORLD = "__world__"
+
+
+@runtime_checkable
+class HasPose(Protocol):
+    """Structural type for any MJCF element that carries a pose."""
+
+    pose: AnyPose
+
+
+def _to_T(posed: HasPose) -> np.ndarray:
+    T = np.eye(4)
+    T[:3, :3] = posed.pose.as_matrix()
+    T[:3, 3] = np.asarray(posed.pose.pos, dtype=float)
+    return T
+
+
+def _from_T(T: np.ndarray) -> PoseQuat:
+    return Quat.from_matrix(T[:3, :3]).as_pose(pos=T[:3, 3])
+
+
+class PoseContext:
+    """
+    Graph of all posed elements in a WorldBody tree for cross-branch pose resolution.
+
+    Every element with a pose (`Body`, any site type, `Frame`, `Camera`, `Light`, etc.) is registered as a node keyed by Python object identity, so no naming is required.
+
+    Attributes:
+        _tm: internal pytransform3d `TransformManager` storing pairwise 4x4 transforms.
+        _registered: set of `id()` values for every element registered during the tree walk.
+
+    """
+
+    def __init__(self, worldbody: WorldBody) -> None:
+        self._tm = TransformManager()
+        self._registered: set[int] = set()
+        self._build(worldbody)
+
+    def _key(self, obj: object) -> str:
+        self._registered.add(id(obj))
+        return str(id(obj))
+
+    def _register_posed_children(self, parent_obj: object, parent_key: str) -> None:
+        """Registers sites, cameras, and lights of a body as child frames."""
+        for attr in ("sites", "cameras", "lights"):
+            for child in getattr(parent_obj, attr, ()):
+                logger.debug(
+                    "registering %s %s", attr[:-1], getattr(child, "name", id(child))
+                )
+                self._tm.add_transform(self._key(child), parent_key, _to_T(child))
+
+    def _walk_frames(self, frm: Frame, parent_key: str) -> None:
+        """Recursively registers a Frame and its nested frames."""
+        logger.debug("registering frame %s", getattr(frm, "name", id(frm)))
+        frm_key = self._key(frm)
+        self._tm.add_transform(frm_key, parent_key, _to_T(frm))
+        for nested in frm.frames:
+            self._walk_frames(nested, frm_key)
+
+    def _build(self, worldbody: WorldBody) -> None:
+        """Walks the full worldbody tree and registers every posed element."""
+        logger.debug("building pose context from worldbody")
+        wb_key = self._key(worldbody)
+        self._tm.add_transform(wb_key, _WORLD, np.eye(4))
+        self._register_posed_children(worldbody, wb_key)
+        for frm in worldbody.frames:
+            self._walk_frames(frm, wb_key)
+        for body in worldbody.bodies:
+            self._walk(body, wb_key)
+
+    def _walk(self, body: Body, parent_key: str) -> None:
+        """Recursively registers a body and all its descendants."""
+        logger.debug("registering body %s", getattr(body, "name", id(body)))
+        body_key = self._key(body)
+        self._tm.add_transform(body_key, parent_key, _to_T(body))
+        self._register_posed_children(body, body_key)
+        for frm in body.frames:
+            self._walk_frames(frm, body_key)
+        for child in body.bodies:
+            self._walk(child, body_key)
+
+    def local_pose(self, frame: HasPose, relative_to: HasPose) -> PoseQuat:
+        """Returns the pose of `frame` expressed in `relative_to`'s coordinate system."""
+        if id(frame) not in self._registered:
+            msg = "frame is not registered in the pose context. Ensure it is part of the worldbody tree before resolving poses."
+            logger.error(msg)
+            raise ValueError(msg)
+        if id(relative_to) not in self._registered:
+            msg = "relative_to is not registered in the pose context. Ensure it is part of the worldbody tree before resolving poses."
+            logger.error(msg)
+            raise ValueError(msg)
+        try:
+            T = self._tm.get_transform(str(id(frame)), str(id(relative_to)))
+        except Exception as exc:
+            msg = f"Cannot resolve transform from frame to relative_to: {exc}"
+            logger.error(msg)
+            raise ValueError(msg) from exc
+        return _from_T(T)
+
+
+@dataclass(frozen=True)
+class PoseRef:
+    """Deferred pose reference resolved via `Mujoco.local_pose`. `relative_to` must be the direct parent body of the element receiving the resolved pose."""
+
+    frame: HasPose = field()
+    """The element whose world-space pose you want to resolve."""
+
+    relative_to: HasPose = field()
+    """The element whose local coordinate system the result is expressed in. Must be the parent body of whatever element receives the resolved pose."""
+
+    def to_quat(self, mjcf: Mujoco) -> PoseQuat:
+        """Resolves this reference to a concrete `PoseQuat` in `relative_to`'s coordinate system."""
+        return mjcf.local_pose(self.frame, self.relative_to)

{mujoco_mojo-2.3.4 → mujoco_mojo-2.3.6}/src/mujoco_mojo/runtime/__init__.py

@@ -1,9 +1,12 @@
 from .load import (
     GeneralLoad,
+    JointFriction,
+    JointLoad,
     Load,
     PointToPointForce,
     ScalarForce,
     ScalarTorque,
+    SiteLoad,
     VectorForce,
     VectorTorque,
 )
@@ -13,12 +16,15 @@ from .video_recorder import VideoRecorder
 
 __all__ = [
     "GeneralLoad",
+    "JointFriction",
+    "JointLoad",
     "Load",
     "PointToPointForce",
     "RuntimeManager",
     "ScalarForce",
     "ScalarTorque",
     "SignalManager",
+    "SiteLoad",
     "VectorForce",
     "VectorTorque",
     "VideoRecorder",