mujoco-toolbox 0.3.4__py3-none-any.whl

mujoco_toolbox/__init__.py

@@ -0,0 +1,77 @@
+"""Mujoco Toolbox.
+==============
+A toolbox for working with MuJoCo simulations. This package provides various utilities and controllers to facilitate
+the simulation process.
+
+Modules
+--------
+- `Wrapper`: Contains the Wrapper class for interfacing with MuJoCo.
+- `Controller`: Includes sineController, cosineController, and randomController for controlling simulations.
+- `Builder`: Contains the Builder class for creating and merging MuJoCo models.
+- `assets`: Contains pre-defined assets for building MuJoCo models.
+
+Constants
+----------
+- CAPTURE_PARAMETERS: List of MjData fields to capture during simulation.
+
+Notes
+-----
+This package is still under development. Report any issues to https://github.com/MGross21/mujoco-toolbox/issues.
+
+"""
+
+from .assets import WORLD_ASSETS, glovebox
+from .builder import Builder
+from .controller import (
+    cosine_controller,
+    live_controller,
+    random_controller,
+    sine_controller,
+    step_controller,
+)
+from .utils import _Platform
+from .wrapper import Wrapper
+
+__version__ = "0.3.4"
+__author__ = "Michael Gross"
+__github_repo__ = "mujoco-toolbox"
+__license__ = "MIT"
+__status__ = "Development"
+__all__ = [
+    "CAPTURE_PARAMETERS",
+    "COMPUTER",
+    "WORLD_ASSETS",
+    "Builder",
+    "Wrapper",
+    "cosine_controller",
+    "glovebox",
+    "live_controller",
+    "random_controller",
+    "sine_controller",
+    "step_controller",
+]
+
+COMPUTER = _Platform()  # Singleton Instance
+MAX_GEOM_SCALAR: int = 2  # Scalar value for mujoco.Renderer.max_geom
+PROGRESS_BAR: bool = False  # Enable/Disable progress bar
+CAPTURE_PARAMETERS = [  # MjData default fields to capture during simulation
+    "time",
+    "qpos",
+    "qvel",
+    "act",
+    "qacc",
+    "xpos",
+    "xquat",
+    "xmat",
+    "ctrl",
+    "sensordata",
+]
+
+if __version__.startswith("0"):
+    from .utils import _print_warning
+    _print_warning(
+        f"{__package__} (v{__version__}) is still under development.",
+        f"Report any issues to https://github.com/MGross21/{__github_repo__}/issues",
+    )
+    del _print_warning
+del _Platform

mujoco_toolbox/assets.py

@@ -0,0 +1,29 @@
+WORLD_ASSETS = """
+  <asset>
+    <texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="32" height="512"/>
+    <texture name="body" type="cube" builtin="flat" mark="cross" width="128" height="128" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1"/>
+    <material name="body" texture="body" texuniform="true" rgba="0.8 0.6 .4 1"/>
+    <texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
+    <material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
+    <texture name="background" type="2d" builtin="flat" width="256" height="256" rgb1="1 1 1"/>
+    <material name="background" texture="background" texuniform="true" rgba="1 1 1 1"/>
+  </asset>
+"""
+
+def glovebox(width=1.25, depth=0.75, height=1.0, glass_thickness=0.05) -> str:
+    return f"""
+<mujoco>
+    <asset>
+        <material name="glass" rgba="1 1 1 0.2"/>
+    </asset>
+    <worldbody>
+        <body name="walls" pos="0 0 0">
+            <geom type="box" size="{width/2} {glass_thickness/2} {height/2}" material="glass" pos="0 {depth/2+glass_thickness/2} {height/2}"/>
+            <geom type="box" size="{width/2} {glass_thickness/2} {height/2}" material="glass" pos="0 {-depth/2-glass_thickness/2} {height/2}"/>
+            <geom type="box" size="{glass_thickness/2} {depth/2} {height/2}" material="glass" pos="{width/2-glass_thickness/2} 0 {height/2}"/>
+            <geom type="box" size="{glass_thickness/2} {depth/2} {height/2}" material="glass" pos="{-width/2+glass_thickness/2} 0 {height/2}"/>
+            <geom type="plane" size="{width/2} {depth/2+glass_thickness} {glass_thickness/2}" material="glass" pos="0 0 {height}"/>
+        </body>
+    </worldbody>
+</mujoco>
+"""

mujoco_toolbox/builder.py

@@ -0,0 +1,229 @@
+import os
+import xml.etree.ElementTree as ET
+
+
+class Builder:
+    """Builder class for loading, merging, and saving MuJoCo models from XML strings or file paths.
+
+    Attributes:
+        xml (str): XML string of the model.
+
+    Examples:
+        >>> # Load a model
+        >>> builder1 = Builder("path/to/model1.xml")
+        >>> # Merge two models from two object instances
+        >>> builder2 = Builder("path/to/model2.xml")
+        >>> merged_builder = builder1 + builder2
+        >>> merged_builder.save("path/to/merged_model.xml")
+        >>> # Merge two models from two file paths
+        >>> builder = Builder("path/to/model1.xml", "path/to/model2.xml")
+        >>> builder.save("path/to/merged_model.xml")
+
+    """
+
+    def __init__(self, *args: str) -> None:
+        """Initialize with an XML string or a file path. Input is required."""
+        for arg in args:
+            if not isinstance(arg, str):
+                msg = "Input must be an XML string or a file path"
+                raise TypeError(msg)
+        if len(args) == 0:
+            msg = "Input is required to initialize the Builder"
+            raise ValueError(msg)
+        if len(args) > 1:
+            # Create a new Builder instance for each additional argument and merge them
+            self.tree, self.root = self._load_model(args[0])
+            for additional_arg in args[1:]:
+                additional_builder = Builder(additional_arg)
+                self.__add__(additional_builder)
+        else:
+            xml_input = args[0]
+            self.tree, self.root = self._load_model(xml_input)
+
+    def _load_model(self, xml_input):
+        """Load a MuJoCo model from a file or XML string."""
+        if not isinstance(xml_input, str):
+            msg = "Input must be either an XML string or a file path"
+            raise ValueError(msg)
+
+        # Check if input is an XML string
+        if xml_input.strip().startswith("<"):
+            # Try parsing as XML string
+            try:
+                root = ET.fromstring(xml_input)
+                tree = ET.ElementTree(root)
+
+                # If it's a URDF (robot tag), wrap it in mujoco
+                if root.tag == "robot":
+                    return self._wrap_urdf_in_mujoco(tree)
+
+                # If root is not mujoco, wrap it
+                if root.tag != "mujoco":
+                    mujoco_elem = ET.Element("mujoco")
+                    mujoco_elem.append(root)
+                    tree = ET.ElementTree(mujoco_elem)
+
+                return tree, tree.getroot()
+            except ET.ParseError:
+                msg = "Invalid XML string provided"
+                raise ValueError(msg)
+        else:
+            # Treat input as file path
+            if not os.path.exists(xml_input):
+                msg = f"File not found: {xml_input}"
+                raise FileNotFoundError(msg)
+
+            # Handle URDF files
+            if xml_input.lower().endswith(".urdf"):
+                return self._wrap_urdf_in_mujoco(ET.parse(xml_input))
+
+            # Parse regular XML file
+            try:
+                tree = ET.parse(xml_input)
+                root = tree.getroot()
+
+                # If root is not mujoco, wrap it
+                if root.tag != "mujoco":
+                    mujoco_elem = ET.Element("mujoco")
+                    mujoco_elem.append(root)
+                    tree = ET.ElementTree(mujoco_elem)
+
+                return tree, tree.getroot()
+            except ET.ParseError:
+                msg = f"Invalid XML file: {xml_input}"
+                raise ValueError(msg)
+
+    def _wrap_urdf_in_mujoco(self, urdf_input):
+        """Wrap a URDF inside <mujoco> tags for MuJoCo compatibility.
+        Input can be a file path, ElementTree, or XML string.
+        """
+        # Handle string input (could be a file path or XML string)
+        if isinstance(urdf_input, str):
+            # Check if it's an XML string
+            if urdf_input.strip().startswith("<"):
+                try:
+                    root = ET.fromstring(urdf_input)
+                    tree = ET.ElementTree(root)
+                except ET.ParseError:
+                    msg = "Invalid URDF XML string provided"
+                    raise ValueError(msg)
+            else:
+                # It's a file path
+                try:
+                    tree = ET.parse(urdf_input)
+                except:
+                    msg = f"Could not parse URDF file: {urdf_input}"
+                    raise ValueError(msg)
+        # Handle ElementTree input
+        elif isinstance(urdf_input, ET.ElementTree):
+            tree = urdf_input
+        else:
+            msg = "URDF input must be a string or ElementTree"
+            raise TypeError(msg)
+
+        root = tree.getroot()
+
+        # Handle the case where the root might already be a mujoco element
+        if root.tag == "mujoco":
+            return tree, root
+
+        # Create mujoco element and append the URDF root
+        mujoco_elem = ET.Element("mujoco")
+        mujoco_elem.append(root)
+
+        # Create a new tree with the <mujoco> root
+        return ET.ElementTree(mujoco_elem), mujoco_elem
+
+    def _merge_tags(self, tag_name, root_1, root_2) -> None:
+        """Merge a specific tag (e.g., worldbody, asset) from two models."""
+        section_1 = root_1.find(tag_name)
+        section_2 = root_2.find(tag_name)
+
+        # If target section doesn't exist in model 1 but exists in model 2, create it
+        if section_1 is None and section_2 is not None:
+            section_1 = ET.SubElement(root_1, tag_name)
+
+        # Only proceed if both sections exist now
+        if section_1 is not None and section_2 is not None:
+            # Merge the contents by appending all elements from model 2 to model 1
+            for element in list(section_2):  # Use list() to avoid modification during iteration
+                section_1.append(element)
+
+    def __add__(self, other):
+        """Implement the + operator for merging two models."""
+        model_root = None
+
+        if isinstance(other, Builder):
+            # Merge the current model with another Builder instance
+            model_root = other.root
+        elif isinstance(other, str):
+            # If the input is a string, load as XML
+            _, model_root = self._load_model(other)
+        else:
+            msg = "Addition only supported between Builder objects or with XML strings/files"
+            raise TypeError(msg)
+
+        # List of sections to merge
+        sections_to_merge = ["asset", "worldbody", "camera", "light", "contact", "equality",
+                            "sensor", "actuator", "default", "tendon", "include"]
+
+        # Merge all relevant sections
+        for section in sections_to_merge:
+            self._merge_tags(section, self.root, model_root)
+
+        # Return self to allow for method chaining
+        return self
+
+    def __radd__(self, other):
+        """Implement the reverse + operator for merging two models."""
+        if other == 0:  # Handle sum() starting value
+            return self
+        return self.__add__(other)
+
+    def save(self, file_path) -> "os.path":
+        """Save the merged model to a file."""
+        if self.tree is not None:
+            # Format the XML with proper indentation before saving
+            self._indent_xml(self.root)
+            self.tree.write(file_path, encoding="utf-8", xml_declaration=True)
+        else:
+            msg = "No model loaded. Cannot save."
+            raise ValueError(msg)
+        return os.path.abspath(file_path)
+
+    def _indent_xml(self, elem, level=0) -> None:
+        """Add proper indentation to make the XML file more readable."""
+        i = "\n" + level * "  "
+        if len(elem):
+            if not elem.text or not elem.text.strip():
+                elem.text = i + "  "
+            if not elem.tail or not elem.tail.strip():
+                elem.tail = i
+            for sub_elem in elem:
+                self._indent_xml(sub_elem, level + 1)
+            if not elem[-1].tail or not elem[-1].tail.strip():
+                elem[-1].tail = i
+        elif level and (not elem.tail or not elem.tail.strip()):
+            elem.tail = i
+
+    def __str__(self) -> str:
+        """Return the XML string of the model."""
+        if self.tree is not None:
+            # Format the XML with proper indentation
+            self._indent_xml(self.root)
+            return ET.tostring(self.root, encoding="unicode", method="xml")
+        msg = "No model loaded. Cannot generate string."
+        raise ValueError(msg)
+
+    def __repr__(self) -> str:
+        """Return the XML string of the model."""
+        return self.__str__()
+
+    def __len__(self) -> int:
+        """Return the number of elements in the model."""
+        return len(self.root) if self.root is not None else 0
+
+    @property
+    def xml(self):
+        """XML string of the model."""
+        return str(self)

mujoco_toolbox/controller.py

@@ -0,0 +1,164 @@
+from numpy import cos, pi, random, sin
+
+
+def _apply_control(model, data, value, joint=None, axis=None, delay=0) -> None:
+    """Common helper function for controller logic to reduce redundancy.
+
+    Args:
+        model: MuJoCo model.
+        data: MuJoCo data.
+        value: Control value to apply.
+        joint (list[int], optional): Joints to apply control to.
+        axis (int, optional): Axis to apply control to.
+        delay (float, optional): Delay before applying control.
+
+    """
+    # Don't apply control until after delay
+    if data.time < delay:
+        return
+
+    # Determine targets if not explicitly provided
+    if joint is None and axis is None:
+        joint = range(model.nu)
+
+    # Apply control to joints or specific axis
+    if joint is not None and model.nu > 0:
+        for j in joint:
+            data.ctrl[j] = value
+    elif axis is not None:
+        data.qpos[axis] = value
+
+
+def step_controller(model, data, **kwargs) -> None:
+    """A step controller for the simulation.
+
+    Args:
+        amplitude (float): The amplitude of the step signal (default=1).
+        joint (list[int]): The joints to apply the step signal to (default=all).
+        axis (int): The axis to apply the step signal to (default=None).
+        delay (float): The delay before applying the step signal (default=0).
+
+    Returns:
+        None
+
+    """
+    amplitude = kwargs.get("amplitude", 1)
+    joint = kwargs.get("joint")
+    axis = kwargs.get("axis")
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        msg = "Delay must be non-negative."
+        raise ValueError(msg)
+    if joint is not None and axis is not None:
+        msg = "Cannot specify both 'joint' and 'axis'."
+        raise ValueError(msg)
+
+    _apply_control(model, data, amplitude, joint=joint, axis=axis, delay=delay)
+
+
+def sine_controller(model, data, **kwargs) -> None:
+    """A simple sine wave controller for the simulation.
+
+    Args:
+        amplitude (float): The amplitude of the sine wave (default=1).
+        frequency (float): The frequency of the sine wave (default=1).
+        phase (float): The phase shift of the sine wave (default=0).
+        joint (list[int]): The joint to apply the sine wave to (default=all).
+        delay (float): The delay before applying the sine wave (default=0).
+
+    Returns:
+        None
+
+    """
+    amplitude = kwargs.get("amplitude", 1)
+    frequency = kwargs.get("frequency", 1)
+    phase = kwargs.get("phase", 0)
+    joint = kwargs.get("joint")
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        msg = "Delay must be non-negative."
+        raise ValueError(msg)
+
+    value = amplitude * sin(2 * pi * frequency * data.time + phase)
+    _apply_control(model, data, value, joint=joint, delay=delay)
+
+
+def cosine_controller(model, data, **kwargs) -> None:
+    """A simple cosine wave controller for the simulation.
+
+    Args:
+        amplitude (float): The amplitude of the cosine wave (default=1).
+        frequency (float): The frequency of the cosine wave (default=1).
+        phase (float): The phase shift of the cosine wave (default=0).
+        joint (list[int]): The joint to apply the cosine wave to (default=all).
+        delay (float): The delay before applying the cosine wave (default=0).
+
+    Returns:
+        None
+
+    """
+    amplitude = kwargs.get("amplitude", 1)
+    frequency = kwargs.get("frequency", 1)
+    phase = kwargs.get("phase", 0)
+    joint = kwargs.get("joint")
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        msg = "Delay must be non-negative."
+        raise ValueError(msg)
+
+    value = amplitude * cos(2 * pi * frequency * data.time + phase)
+    _apply_control(model, data, value, joint=joint, delay=delay)
+
+
+def random_controller(model, data, **kwargs) -> None:
+    """A random controller for the simulation.
+
+    Args:
+        amplitude (float): The maximum amplitude of the random signal (default=1).
+        joint (list[int]): The joints to apply the random signal to (default=all).
+        axis (int): The axis to apply the random signal to (default=None).
+        delay (float): The delay before applying the random signal (default=0).
+
+    Returns:
+        None
+
+    """
+    amplitude = kwargs.get("amplitude", 1)
+    joint = kwargs.get("joint")
+    axis = kwargs.get("axis")
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        msg = "Delay must be non-negative."
+        raise ValueError(msg)
+    if joint is not None and axis is not None:
+        msg = "Cannot specify both 'joint' and 'axis'."
+        raise ValueError(msg)
+
+    value = amplitude * random.rand()
+    _apply_control(model, data, value, joint=joint, axis=axis, delay=delay)
+
+def live_controller(model, data, **kwargs) -> None:
+    """A real-time controller for the simulation.
+
+    Args:
+        controller_params (dict): Dictionary of parameters to pass to the controller.
+
+    Returns:
+        None
+
+    """
+    from .utils import _print_warning
+    for key, value in kwargs.get("controller_params", {}).items():
+        if hasattr(data, key):
+            setattr(data, key, value)
+        else:
+            _print_warning(f"'{key}' is not a valid attribute of MjData. Skipping...")
+
+    # untested
+    # if hasattr(data, 'control'):
+    #     data.control = kwargs.get("control", data.control)
+

mujoco_toolbox/utils.py

@@ -0,0 +1,134 @@
+import platform
+import sys
+from functools import wraps
+from multiprocessing import cpu_count
+from time import time
+
+from colorama import Fore, Style, init
+from screeninfo import get_monitors
+
+
+def timer(func):
+    @wraps(func)
+    def time_wrapper(*args, **kwargs):
+        start_time = time()
+        result = func(*args, **kwargs)
+        end_time = time()
+        total_time = end_time - start_time
+        print(f"Function {func.__name__} took {total_time:.4f} seconds to execute")
+        return result
+
+    return time_wrapper
+
+def _print_warning(*args, prefix: bool = True) -> None:
+    """Prints the first argument as an orange warning message and prints the other arguments in regular text, all on the same line.
+
+    Args:
+        *args: First argument is printed as an orange warning, other arguments are printed normally.
+
+    """
+    if args:
+        if prefix:
+            # Make sure the first argument is printed with the 'WARNING:' prefix and in orange
+            message = f"WARNING: {args[0]}"
+        else:
+            message = args[0]
+
+        # Check if we're on Windows or a Unix-like system for color support
+        if sys.platform.startswith("win"):
+            try:
+                init()  # Initialize colorama
+                print(f"{Fore.YELLOW}{message}{Style.RESET_ALL}", end=" ")  # Yellow color for warning
+                print(*args[1:])
+            except ImportError:
+                # If colorama is not available, print without color
+                print(message, end=" ")
+                print(*args[1:])
+        else:
+            # ANSI escape code for orange (for Unix-like systems)
+            print(f"\033[38;5;214m{message}\033[0m", end=" ")  # Orange ANSI escape code
+            print(*args[1:])
+
+def _print_success(*args, prefix: bool = True) -> None:
+    """Prints the first argument as a green success message and prints the other arguments in regular text, all on the same line.
+
+    Args:
+        *args: First argument is printed as a green success message, other arguments are printed normally.
+
+    """
+    if args:
+        if prefix:
+            # Make sure the first argument is printed with the 'SUCCESS:' prefix and in green
+            message = f"SUCCESS: {args[0]}"
+        else:
+            message = args[0]
+
+        # Check if we're on Windows or a Unix-like system for color support
+        if sys.platform.startswith("win"):
+            try:
+                init()  # Initialize colorama
+                print(f"{Fore.GREEN}{message}{Style.RESET_ALL}", end=" ")  # Green color for success
+                print(*args[1:])
+            except ImportError:
+                # If colorama is not available, print without color
+                print(message, end=" ")
+                print(*args[1:])
+        else:
+            # ANSI escape code for green (for Unix-like systems)
+            print(f"\033[32m{message}\033[0m", end=" ")  # Green ANSI escape code
+            print(*args[1:])
+
+class _Platform:
+    def __init__(self) -> None:
+        self.OS = platform.system()
+        self.PROCESSOR = platform.processor()
+        self.MACHINE = platform.machine()
+        self.PYTHON_VERSION = platform.python_version()
+        self.IDE = self.get_developer_mode()
+        self.CPU_COUNT = cpu_count()
+        self.ARCHITECTURE = platform.architecture()[0]
+        self.PYTHON_IMPLEMENTATION = platform.python_implementation()
+        self.RESOLUTION, self.GUI_ENABLED = self.get_resolution()
+        # self.NUM_MONITORS = get_monitors().__len__()
+
+    def __repr__(self) -> str:
+        return (
+            f"Operating System: {self.OS}\n"
+            f"Architecture: {self.ARCHITECTURE}\n"
+            f"Processor: {self.PROCESSOR}\n"
+            f"Machine Type: {self.MACHINE}\n"
+            f"CPU Count: {self.CPU_COUNT}\n"
+            f"Python Version: {self.PYTHON_VERSION}\n"
+            f"IDE: {self.IDE}\n"
+            f"Python Implementation: {self.PYTHON_IMPLEMENTATION}\n"
+            f"Resolution (disp0): {self.RESOLUTION}\n"
+            f"GUI Enabled: {self.GUI_ENABLED}"
+        )
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+    @staticmethod
+    def get_developer_mode() -> str:
+        if "ipykernel" in sys.modules and not sys.stdin.isatty():
+            return "jupyter"
+        return "terminal" if hasattr(sys, "ps1") else "script"
+
+    @staticmethod
+    def get_resolution():
+        """Detects screen resolution and potential headless operation."""
+        try:
+            monitor = get_monitors()[0]
+            return (monitor.width, monitor.height), True
+        except Exception:
+            _print_warning("Detected headless operation. Disabling GUI...")
+            return (1920, 1080), False
+
+class _EmptyContextManager:
+    __slots__ = ()
+    def __init__(self, *args, **kwargs) -> None:
+        pass
+    def __enter__(self):
+        pass
+    def __exit__(self, _exc_type, _exc_value, _traceback):
+        pass