Skip to content

Mojo model

mojo_model

UserData

Bases: MojoBaseModel

Base class for all user-defined simulation data.

MojoModel

Bases: MojoBaseModel, StochasBase

Mojo is the highest level watcher which manages model definitions.

mjcf class-attribute instance-attribute

Python
mjcf: Mujoco = Field(default_factory=Mujoco)

MuJoCo MJCF model to be written to XML.

user_data class-attribute instance-attribute

Python
user_data: SerializeAsAny[UserData] | None = None

User defined data for the model. This is used for transferring information from one function to another (generator to runtime or objective function).

trial_dir property

Python
trial_dir: Path

Absolute path to this trial's workspace directory.

Set automatically after the generator runs. Available in runtime and objective functions. Raises RuntimeError if accessed inside the generator, where the directory has not yet been created.

get_user_data

Python
get_user_data(cls: type[T]) -> T

Returns the user_data re-validated into the requested class.

Source code in src/mujoco_mojo/mojo_model.py
Python
def get_user_data(self, cls: type[T]) -> T:
    """Returns the user_data re-validated into the requested class."""
    if self.user_data is None:
        msg = "Unable to get user_data since it is None."
        logger.error(msg)
        raise ValueError(msg)

    if not isinstance(self.user_data, cls):
        data = self.user_data
        if isinstance(data, BaseModel):
            data = data.model_dump()

        self.user_data = cls.model_validate(data)

    return self.user_data

clear_unpickleable_data

Python
clear_unpickleable_data() -> None

Clear unpickleable objects from the model tree before serialization.

This walks through all geoms in the MJCF tree and clears cached collision managers and proximity queries that contain Cython objects which cannot be pickled for multiprocessing.

Source code in src/mujoco_mojo/mojo_model.py
Python
def clear_unpickleable_data(self) -> None:
    """
    Clear unpickleable objects from the model tree before serialization.

    This walks through all geoms in the MJCF tree and clears cached collision managers and proximity queries that contain Cython objects which cannot be pickled for multiprocessing.
    """
    from mujoco_mojo.utils.proximity_mixin import ProximityMixin

    if self.mjcf.worldbody is None:
        return

    bodies = self.mjcf.worldbody.walk_bodies(include_self=True)
    for body in bodies:
        for geom in body.geoms:
            if isinstance(geom, ProximityMixin):
                geom.clear_unpickleable()