base

BaseCompositionElement

Bases: ABC

Base class for composable elements in the MuJoCo model, providing common functionality such as compiling to MuJoCo model/data and exporting.

Source code in src/flygym/compose/base.py

class BaseCompositionElement(ABC):
    """Base class for composable elements in the MuJoCo model, providing common
    functionality such as compiling to MuJoCo model/data and exporting."""

    @property
    @abstractmethod
    def mjcf_root(self) -> mjcf.RootElement:
        """The root MJCF element of this composition element."""
        pass

    def compile(self) -> tuple[mj.MjModel, mj.MjData]:
        """Compile the MJCF model into MuJoCo MjModel and MjData objects. This is where
        `dm_control.mjcf` "hands off" the model that it composes to `mujoco` for
        simulation. Things like the ordering of generalized coordinates (qpos) are
        determined here."""
        physics = mjcf.Physics.from_mjcf_model(self.mjcf_root)
        return physics.model._model, physics.data._data

    def save_xml_with_assets(
        self, output_dir: PathLike, xml_filename: str = None
    ) -> None:
        """Export the MJCF model to a directory, including a XML file (filename defaults
        to the model name if not specified) and all associated assets (e.g. meshes)."""
        Path(output_dir).mkdir(parents=True, exist_ok=True)
        mjcf.export_with_assets(self.mjcf_root, str(output_dir), xml_filename)

mjcf_root abstractmethod property

The root MJCF element of this composition element.

compile()

Compile the MJCF model into MuJoCo MjModel and MjData objects. This is where dm_control.mjcf "hands off" the model that it composes to mujoco for simulation. Things like the ordering of generalized coordinates (qpos) are determined here.

Source code in src/flygym/compose/base.py

def compile(self) -> tuple[mj.MjModel, mj.MjData]:
    """Compile the MJCF model into MuJoCo MjModel and MjData objects. This is where
    `dm_control.mjcf` "hands off" the model that it composes to `mujoco` for
    simulation. Things like the ordering of generalized coordinates (qpos) are
    determined here."""
    physics = mjcf.Physics.from_mjcf_model(self.mjcf_root)
    return physics.model._model, physics.data._data

save_xml_with_assets(output_dir, xml_filename=None)

Export the MJCF model to a directory, including a XML file (filename defaults to the model name if not specified) and all associated assets (e.g. meshes).

Source code in src/flygym/compose/base.py

def save_xml_with_assets(
    self, output_dir: PathLike, xml_filename: str = None
) -> None:
    """Export the MJCF model to a directory, including a XML file (filename defaults
    to the model name if not specified) and all associated assets (e.g. meshes)."""
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    mjcf.export_with_assets(self.mjcf_root, str(output_dir), xml_filename)