mujoco-toolbox 0.1.8__py3-none-any.whl

mujoco_toolbox/Controller.py

@@ -0,0 +1,131 @@
+from numpy import cos, pi, random, sin
+
+
+def _apply_control(model, data, value, joint=None, axis=None, delay=0):
+    """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 stepController(model, data, **kwargs):
+    """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", None)
+    axis = kwargs.get("axis", None)
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        raise ValueError("Delay must be non-negative.")
+    if joint is not None and axis is not None:
+        raise ValueError("Cannot specify both 'joint' and 'axis'.")
+
+    _apply_control(model, data, amplitude, joint=joint, axis=axis, delay=delay)
+
+
+def sineController(model, data, **kwargs):
+    """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", None)
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        raise ValueError("Delay must be non-negative.")
+
+    value = amplitude * sin(2 * pi * frequency * data.time + phase)
+    _apply_control(model, data, value, joint=joint, delay=delay)
+
+
+def cosineController(model, data, **kwargs):
+    """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", None)
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        raise ValueError("Delay must be non-negative.")
+
+    value = amplitude * cos(2 * pi * frequency * data.time + phase)
+    _apply_control(model, data, value, joint=joint, delay=delay)
+
+
+def randomController(model, data, **kwargs):
+    """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", None)
+    axis = kwargs.get("axis", None)
+    delay = kwargs.get("delay", 0)
+
+    if delay < 0:
+        raise ValueError("Delay must be non-negative.")
+    if joint is not None and axis is not None:
+        raise ValueError("Cannot specify both 'joint' and 'axis'.")
+
+    value = amplitude * random.rand()
+    _apply_control(model, data, value, joint=joint, axis=axis, delay=delay)

mujoco_toolbox/Utils.py

@@ -0,0 +1,53 @@
+from functools import wraps
+from sys import platform
+from time import time
+
+
+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
+
+        from . import VERBOSITY
+
+        if VERBOSITY:
+            print(f"Function '{func.__name__}' took {total_time:.4f} seconds")
+        return result
+
+    return time_wrapper
+
+
+def print_warning(*args) -> 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:
+        # Make sure the first argument is printed with the 'WARNING:' prefix and in orange
+        message = f"WARNING: {args[0]}"
+
+        # Check if we're on Windows or a Unix-like system for color support
+        if platform.startswith("win"):
+            try:
+                from colorama import Fore, Style, init
+
+                init()  # Initialize colorama
+                print(
+                    f"{Fore.YELLOW}{message}{Style.RESET_ALL}", end=" "
+                )  # Yellow color for warning
+                print(
+                    *args[1:], end=""
+                )  # No newline, print all args separated by spaces
+            except ImportError:
+                # If colorama is not available, print without color
+                print(message, end=" ")
+                print(*args[1:], end="")
+        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:], end="")  # No newline, print all args separated by spaces

mujoco_toolbox/Wrapper.py

@@ -0,0 +1,741 @@
+import os
+import sys
+import xml.etree.ElementTree as ET
+from collections import defaultdict
+from functools import lru_cache
+from typing import Any, Dict, List, Optional, Tuple
+
+import matplotlib.pyplot as plt
+import mediapy as media
+import mujoco
+import numpy as np
+import trimesh
+import yaml
+from screeninfo import get_monitors
+from tqdm import tqdm as barTerminal
+from tqdm.notebook import tqdm as barNotebook
+from screeninfo import get_monitors
+import trimesh
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple, Union, Optional, Callable, Any, TypeAlias
+
+from .Utils import print_warning, timer
+
+assert sys.version_info >= (3, 10), "This code requires Python 3.10.0 or later."
+assert mujoco.__version__ >= "2.0.0", "This code requires MuJoCo 2.0.0 or later."
+
+mjModel: TypeAlias = mujoco.MjModel
+mjData: TypeAlias = mujoco.MjData
+
+
+class Wrapper(object):
+    """A class to handle MuJoCo simulations and data capture."""
+
+    def __init__(self, xml:str, duration:int=10, fps:int=30, resolution:Tuple[int,int]=(400,300), initialConditions:Dict[str, List]={}, controller:Optional[Callable[[mjModel, mjData, Any], None]]=None, *args, **kwargs):
+        self._load_model(xml, **kwargs)
+
+        self.duration = duration
+        self.fps = fps
+        self.resolution = resolution  # recursively sets width and height
+        self.init_conditions = initialConditions
+        self.controller = controller
+
+        # Predefined simulation parameters but can be overridden
+        # TODO: Currently Causing Bugs when occluded from XML Code
+        self.ts = kwargs.get("ts", self._model.opt.timestep)
+        self.gravity = kwargs.get("gravity", self._model.opt.gravity)
+
+        self._data = mujoco.MjData(self._model)
+
+        # Auto-Populate the names of bodies, joints, and actuators
+        self._body_names = [self._model.body(i).name for i in range(self._model.nbody)]
+        self._geom_names = [self._model.geom(i).name for i in range(self._model.ngeom)]
+        self._joint_names = [self._model.joint(i).name for i in range(self._model.njnt)]
+        self._actuator_names = [
+            self._model.actuator(i).name for i in range(self._model.nu)
+        ]
+
+    def _load_model(self, xml: str, **kwargs: Any) -> None:
+        """Load a MuJoCo model from an XML file or a string."""
+        try:
+            # Convert the XML path to an absolute path
+            xml_path = os.path.abspath(xml)
+
+            # Check if the path exists
+            if os.path.exists(xml_path):
+                # Extract and validate file extension
+                extension = os.path.splitext(xml_path)[1].lower()[1:]
+
+                if extension == "xml":
+                    self._load_xml_file(xml_path)
+                elif extension == "urdf":
+                    self._load_urdf_file(xml_path, **kwargs)
+                else:
+                    raise ValueError(
+                        f"Unsupported file extension: '{extension}'. Please provide an XML or URDF file."
+                    )
+            else:
+                # If the file doesn't exist, assume it's a string and attempt to load it as XML
+                if "<mujoco>" in xml or "<robot>" in xml:
+                    self._load_xml_string(xml)
+                else:
+                    raise FileNotFoundError(
+                        f"Model file not found: {xml_path}. Ensure the file path is correct and accessible."
+                    )
+
+        except FileNotFoundError as e:
+            raise FileNotFoundError(f"Failed to load the MuJoCo model: {e}") from e
+
+        except ValueError as e:
+            raise ValueError(
+                f"Invalid value encountered while loading the model: {e}"
+            ) from e
+
+        except Exception as e:
+            raise Exception(
+                f"Unexpected error while loading the MuJoCo model: {e}"
+            ) from e
+
+    def _load_xml_file(self, xml: str, **kwargs) -> None:
+        """Load a MuJoCo model from an XML file."""
+        self._model = mujoco.MjModel.from_xml_path(xml)
+        with open(xml, "r") as f:
+            self.xml = f.read()
+        template = kwargs.get('template', None)
+        if template:
+            try:
+                self.xml = self.xml.format(**template)
+            except KeyError as e:
+                raise ValueError(f"Template key error. Ensure the template keys match the placeholders in the XML.") from e
+            except Exception as e:
+                raise ValueError(f"Error formatting XML with template") from e
+        self._model = mujoco.MjModel.from_xml_string(self.xml)
+
+    def _load_urdf_file(self, urdf_path: str, **kwargs: Any) -> None:
+        """Process and load a URDF file for use with MuJoCo."""
+
+        def convert_dae_to_stl(meshdir: str) -> None:
+            """Convert all DAE files in a directory (including subdirectories) to STL."""
+            if not os.path.exists(meshdir):
+                raise FileNotFoundError(f"Directory not found: {meshdir}")
+            for filename in os.listdir(meshdir):
+                if filename.lower().endswith(".dae"):
+                    dae_path = os.path.join(meshdir, filename)
+                    stl_path = os.path.splitext(dae_path)[0] + ".stl"
+                    try:
+                        trimesh.load_mesh(dae_path).export(stl_path)
+                        print(f"Converted: {os.path.basename(dae_path)} -> {os.path.basename(stl_path)}")
+                    except Exception as e:
+                        raise ValueError(f"Error converting {filename}") from e
+                        
+        try:
+            from . import VERBOSITY
+            robot = ET.parse(urdf_path).getroot()
+            mujoco_tag = ET.Element("mujoco")
+            
+            # Get main meshdir (parent directory for meshes)
+            meshdir = kwargs.get("meshdir", "meshes/")  # Default as tuple
+            if not os.path.isabs(meshdir):
+                urdf_dir = os.path.dirname(os.path.abspath(urdf_path))
+                meshdir = os.path.join(urdf_dir, meshdir)
+
+            # Ensure meshdir exists
+            if not os.path.exists(meshdir):
+                raise FileNotFoundError(f"Mesh directory not found: {meshdir}")
+            
+            # If no explicit subdirs are provided, auto-detect subdirectories with STL files
+            subdirs = kwargs.get('meshdir_sub', None)
+            if not subdirs:
+                subdirs = [os.path.relpath(root, meshdir) for root, _, files in os.walk(meshdir) if any(f.endswith('.stl') for f in files)]
+                if VERBOSITY:
+                    print(f"Auto-detected subdirectories: {subdirs}")
+
+            # Convert relative subdir paths to absolute based on meshdir
+            full_meshdirs = [
+                os.path.join(meshdir, subdir) if not os.path.isabs(subdir) else subdir
+                for subdir in subdirs
+            ]
+
+            # Add <compiler> tag with the main meshdir (required for MuJoCo)
+            compiler_attrs = {
+                "meshdir": meshdir,  # Use the main meshdir
+                "balanceinertia": kwargs.get("balanceinertia", "true"),
+                "discardvisual": "false",
+            }
+            ET.SubElement(mujoco_tag, "compiler", **compiler_attrs)
+
+            # Create <asset> tag (required for MuJoCo)
+            asset_tag = ET.SubElement(mujoco_tag, "asset")
+
+
+            # Process all valid mesh directories
+            for full_meshdir in full_meshdirs:
+                if os.path.exists(full_meshdir):
+                    convert_dae_to_stl(full_meshdir)  # Convert DAE to STL in the origin directory
+
+                    # Walk through all files in the directory
+                    for root, _, files in os.walk(full_meshdir):
+                        for filename in files:
+                            if filename.lower().endswith('.stl'):
+                                relative_dir = os.path.relpath(root, meshdir)  # Extract subdir name
+                                mesh_name = f"{os.path.splitext(filename)[0]}_{relative_dir.replace(os.sep, '_')}"  # Format: base_visual
+                                mesh_file_relative = os.path.join(relative_dir, filename)
+                                ET.SubElement(asset_tag, "mesh", name=mesh_name, file=mesh_file_relative)
+                    
+            robot.insert(0, mujoco_tag)
+            self.xml = ET.tostring(robot, encoding="unicode").replace('.dae', '.stl')
+            self._model = mujoco.MjModel.from_xml_string(self.xml)
+        except Exception as e:
+            raise ValueError(f"Failed to process URDF file: {e}")
+
+    def _load_xml_string(self, xml: str) -> None:
+        """Load a MuJoCo model from an XML string."""
+        try:
+            self._model = mujoco.MjModel.from_xml_string(xml)
+            self.xml = xml
+        except Exception as e:
+            raise ValueError(f"Failed to load the MuJoCo model from XML string: {e}") from e
+
+    def __str__(self):
+        return self._model.__str__()
+
+    def __repr__(self):
+        return (
+            f"{self.__class__.__name__}(\n"
+            f"  Duration: {self.duration} [fps={self.fps}, ts={self.ts:.0e}]\n"
+            f"  Gravity: {self.gravity},\n"
+            f"  Resolution: {self._width}x{self._height}\n"
+            f"  Bodies ({self.model.nbody}): {', '.join(self._body_names[:5])}{' ...' if len(self._body_names) > 5 else ''}\n"
+            f"  Joints ({self.model.njnt}): {', '.join(self._joint_names[:5])}{' ...' if len(self._joint_names) > 5 else ''}\n"
+            f"  Actuators ({self.model.nu}): {', '.join(self._actuator_names[:5])}{' ...' if len(self._actuator_names) > 5 else ''}\n"
+            f"  Controller: {self.controller.__name__ if self.controller else None}\n" # Returns str name of the function
+            f")"
+        )
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        mujoco.set_mjcb_control(None)
+        return
+
+    @property
+    def model(self) -> mujoco.MjModel:
+        """Read-only property to access the MjModel object."""
+        return self._model
+
+    @property
+    def data(self) -> mujoco.MjData:
+        """Read-only property to access the MjData single-step object."""
+        return self._data
+
+    @property
+    def captured_data(self) -> Dict[str, np.ndarray]:
+        """Read-only property to access the entire captured simulation data."""
+        if self._captured_data is None:
+            raise ValueError("No simulation data captured yet.")
+        return self._captured_data.unwrap()
+
+    @captured_data.deleter
+    def captured_data(self):
+        self._captured_data = None
+
+    @property
+    def frames(self) -> List[np.ndarray]:
+        """Read-only property to access the captured frames."""
+        if self._frames is None:
+            raise ValueError("No frames captured yet.")
+        return self._frames
+
+    @frames.deleter
+    def frames(self):
+        self._frames.clear()
+        import gc
+
+        gc.collect()
+
+    @property
+    def duration(self) -> float:
+        return self._duration
+
+    @duration.setter
+    def duration(self, value):
+        if value < 0:
+            raise ValueError("Duration must be greater than zero.")
+        self._duration = value
+
+    @property
+    def fps(self) -> float:
+        return self._fps
+
+    @fps.setter
+    def fps(self, value):
+        if value < 0:
+            raise ValueError("FPS must be greater than zero.")
+        self._fps = value
+
+    @property
+    def resolution(self) -> Tuple[int, int]:
+        return self._resolution
+
+    @resolution.setter
+    def resolution(self, values):
+        if len(values) != 2:
+            raise ValueError("Resolution must be a tuple of width and height.")
+        if values[0] < 1 or values[1] < 1:
+            raise ValueError("Resolution must be at least 1x1 pixels.")
+
+        try:
+            monitor = get_monitors()[0]
+            screen_width, screen_height = monitor.width, monitor.height
+        except Exception as e:
+            screen_width, screen_height = 1920, 1080
+            print_warning(
+                f"Failed to get screen resolution: {e}. Defaulting to 1920x1080."
+            )
+
+        if values[0] > screen_width or values[1] > screen_height:
+            raise ValueError("Resolution must be less than the screen resolution.")
+
+        self._resolution = tuple(int(value) for value in values)
+        self._width, self._height = self._resolution
+        # Match changes to the model's visual settings
+        self._model.vis.Global.offwidth = self._width
+        self._model.vis.Global.offheight = self._height
+
+    @property
+    def init_conditions(self):
+        return self._initcond
+
+    @init_conditions.setter
+    def init_conditions(self, values):
+        if not isinstance(values, dict):
+            raise ValueError("Initial conditions must be a dictionary.")
+        invalid_keys = [
+            key for key in values.keys() if not hasattr(mujoco.MjData(self._model), key)
+        ]
+        if invalid_keys:
+            valid_keys = SimulationData._get_public_keys(self._data)
+            print(f"Valid initial condition attributes: {', '.join(valid_keys)}")
+            raise ValueError(
+                f"Invalid initial condition attributes: {', '.join(invalid_keys)}"
+            )
+        self._initcond = values
+
+    @property
+    def controller(self) -> Optional[Callable[[mjModel, mjData, Any], None]]:
+        """Controller Function"""
+        return self._controller
+
+    @controller.setter
+    def controller(self, func: Callable[[mjModel, mjData, Any], None]) -> None:
+        if func is not None and not callable(func):
+            raise ValueError("Controller must be a callable function.")
+        self._controller = func
+
+    @property
+    def ts(self) -> float:
+        return self._model.opt.timestep
+
+    @ts.setter
+    def ts(self, value) -> None:
+        if value <= 0:
+            raise ValueError("Timestep must be greater than 0.")
+        self._model.opt.timestep = value
+
+    @property
+    def data_rate(self) -> int:
+        if self._dr is None:
+            raise ValueError(
+                f"Use '{self.runSim.__name__}' first in order to access this value."
+            )
+        return self._dr
+    
+    @data_rate.setter
+    def data_rate(self, value) -> None:
+        if value.is_numeric() and not isinstance(value, int):
+            value = round(value)
+            print_warning(f"Data rate must be an integer. Rounding to the nearest integer ({value}).")
+        if value <= 0:
+            raise ValueError("Data rate must be greater than 0.")
+        # TODO: Check the math on this validation
+        # max_ = int(self._duration / self.ts)
+        # if value > max_:
+        #     print_warning(f"Data rate exceeds the simulation steps. Setting to the maximum possible value ({max_}).")
+        #     value = max_
+        self._dr = value
+
+    @property
+    def gravity(self):
+        return self._model.opt.gravity
+
+    @gravity.setter
+    def gravity(self, values):
+        if len(values) != 3:
+            raise ValueError("Gravity must be a 3D vector.")
+        self._model.opt.gravity = values
+
+    def _setInitialConditions(self):
+        for key, value in self._initcond.items():
+            if hasattr(self._data, key):
+                setattr(self._data, key, value)
+            else:
+                print_warning(f"'{key}' is not a valid attribute of MjData.")
+
+    def _resetSimulation(self):
+        mujoco.mj_resetData(self._model, self._data)
+        self._setInitialConditions()
+        self._frames = []
+        self._captured_data = SimulationData()
+
+    @timer
+    def runSim(self, render=False, camera=None, data_rate=100, multi_thread=False):
+        """Run the simulation with optional rendering and controlled data capture.
+
+        Args:
+            render (bool): If True, renders the simulation.
+            camera (str): The camera view to render from, defaults to None.
+            data_rate (int): How often to capture data, expressed as frames per second.
+
+        Returns:
+            self: The current Wrapper object for method chaining.
+        """
+        try:
+            mujoco.set_mjcb_control(self._controller) if self._controller else None
+            self._resetSimulation()
+            total_steps = int(self._duration / self.ts)
+
+            # Cache frequently used functions and objects for performance
+            mj_step = mujoco.mj_step
+            m = self._model
+            d = self._data
+
+            self._dr = data_rate
+            capture_rate = data_rate * self.ts
+            capture_interval = max(1, int(1.0 / capture_rate))
+            render_interval = max(1, int(1.0 / (self._fps * self.ts)))
+
+            import __main__ as main
+
+            if hasattr(main, "__file__"):
+                PBar = barTerminal
+            else:
+                PBar = barNotebook
+
+            if multi_thread:
+                #    num_threads =  os.cpu_count()
+                # TODO: Implement multi-threading
+                print(
+                    "Multi-threading not yet implemented. Running simulation in single-thread mode."
+                )
+
+            with (
+                PBar(
+                    total=total_steps, desc="Simulation", unit=" step", leave=False
+                ) as pbar,
+                mujoco.Renderer(self._model, self._height, self._width) as renderer,
+            ):
+                step = 0
+                while d.time < self._duration:
+                    mj_step(m, d)
+
+                    # Capture data at the specified rate
+                    if step % capture_interval == 0:
+                        self._captured_data.capture(d)
+
+                    if render and step % render_interval == 0:
+                        (
+                            renderer.update_scene(d)
+                            if camera is None
+                            else renderer.update_scene(d, camera)
+                        )
+                        self._frames.append(renderer.render().copy())
+
+                    pbar.update(1)
+                    step += 1
+
+                    # if verbose:
+                    #     for warning in mujoco.mjtWarning:
+                    #         if d.warning[warning].number > 0:
+                    #             print_warning(f"{warning.name} - {d.warning[warning].number} occurrences")
+                    # else:
+                    #     if any(d.warning[warning].number > 0 for warning in mujoco.mjtWarning):
+                    #         print_warning("Please check MUJOCO_LOG.txt for more details.")
+
+        except Exception as e:
+            print(f"Simulation error: {e}")
+        finally:
+            mujoco.set_mjcb_control(None)
+
+        return self
+
+    def renderFrame(self, t=0, frame=0, title=None) -> Optional[str]:
+        """Render a specific frame as an image.
+
+        Args:
+            t (float): Time in seconds for which the frame should be rendered.
+            frame (int): The frame index to render.
+            title (str): The title of the rendered frame.
+
+        Returns:
+            None
+        """
+        if not self._frames:
+            raise ValueError("No frames captured to render.")
+
+        if t < 0 or frame < 0:
+            raise ValueError(
+                "Time and frame index must be greater than or equal to zero."
+            )
+
+        if frame and t:
+            raise ValueError("Can only specify singular time or frame parameter")
+
+        try:
+            if t > 0:
+                frame = self.t2f(t)  # Convert time to frame index
+            else:
+                frame = int(frame)
+
+            plt.imshow(self._frames[frame])
+            plt.axis("off")
+            plt.title(title or f"Frame {frame}", loc="center")
+            plt.show()
+
+        except IndexError as e:
+            raise ValueError(f"Invalid frame index: {e}")
+        except TypeError as e:
+            raise ValueError(f"Invalid type for time or frame: {e}")
+        except Exception as e:
+            raise RuntimeError(f"Unexpected error while rendering frame: {e}")
+
+    def renderMedia(self, codec="gif", title=None, save=False) -> media:
+        """Render the simulation as a video or GIF, with an option to save to a file.
+
+        Args:
+            codec (str): The media format to use ("gif" or "mp4").
+            title (str): The filename or window title for the media.
+            save (bool): Whether to save the media to a file.
+        """
+        if not self._frames:
+            raise ValueError("No frames captured to create media.")
+        # Display the media in a window
+        if not save:
+            media.show_video(
+                self._frames,
+                fps=self._fps,
+                width=self._width,
+                height=self._height,
+                codec=codec,
+                title=title,
+            )
+            return
+
+        if not title.endswith(f".{codec}"):
+            title += f".{codec}"
+
+        # Save the frames to the specified file
+        available_codecs = ["gif", "mp4", "h264", "hevc", "vp9"]
+        if codec in available_codecs:
+            media.write_video(
+                title if not None else "render",
+                self._frames,
+                fps=self._fps,
+                codec=codec,
+            )
+        else:
+            raise ValueError(
+                f"Unsupported codec '{codec}'. Supported codecs are {', '.join(available_codecs)}"
+            )
+
+        path = os.path.abspath(title)
+        print(f"Media saved to {path}")
+        return path
+
+    @lru_cache(maxsize=100)
+    def t2f(self, t: float) -> int:
+        """Convert time to frame index."""
+        return min(
+            int(t * self._fps), int(self._duration * self._fps) - 1
+        )  # Subtract 1 to convert to 0-based index
+
+    @lru_cache(maxsize=100)
+    def f2t(self, frame: int) -> float:
+        """Convert frame index to time."""
+        return frame / self._fps
+
+    def getBodyData(self, body_name: str, data_name: str = None) -> np.ndarray:
+        """Get the data for a specific body in the simulation.
+
+        Args:
+            body_name (str): The name of the body to retrieve data for.
+            data_name (str): The name of the data to retrieve.
+
+        Returns:
+            np.ndarray: The data for the specified body.
+        """
+        if body_name not in self._body_names:
+            raise ValueError(f"Body '{body_name}' not found in the model.")
+        else:
+            body_id = mujoco.mj_name2id(
+                self._model, mujoco.mjtObj.mjOBJ_BODY, body_name
+            )
+
+        if data_name is None:
+            return self._captured_data.unwrap()[body_id]
+        if data_name not in self._captured_data.unwrap():
+            raise ValueError(f"Data '{data_name}' not found for body '{body_name}'.")
+        return self._captured_data.unwrap()[body_id][data_name]
+
+    def getID(self, id: int) -> str:
+        """Get the name of a body given its ID.
+
+        Args:
+            id (int): The ID of the body.
+
+        Returns:
+            str: The name of the body.
+        """
+        if id < 0 or id >= self._model.nbody:
+            raise ValueError(f"Invalid body ID: {id}")
+        return mujoco.mj_id2name(self._model, mujoco.mjtObj.mjOBJ_BODY, id)
+
+    def saveYAML(self, name="Model"):
+        """Save simulation data to a YAML file.
+
+        Args:
+            name (str): The filename for the YAML file.
+
+        Returns:
+            None
+        """
+        if not name.endswith(".yml"):
+            name += ".yml"
+
+        try:
+            # Convert simData's NumPy arrays or lists to a YAML-friendly format
+            serialized_data = {
+                k: (v.tolist() if isinstance(v, np.ndarray) else v)
+                for k, v in self.captured_data.items()
+            }
+
+            with open(name, "w") as f:
+                yaml.dump(serialized_data, f, default_flow_style=False)
+
+            print(f"Simulation data saved to {name}")
+        except Exception as e:
+            print(f"Failed to save data to YAML: {e}")
+
+
+class SimulationData(object):
+    """A class to store and manage simulation data."""
+
+    __slots__ = ['_d']
+
+    def __init__(self):
+        self._d = defaultdict(list)
+
+    def capture(self, mj_data: mujoco.MjData):
+        """Capture data from MjData object, storing specified or all public simulation data."""
+        from . import CAPTURE_PARAMETERS, VERBOSITY
+        keys = self._get_public_keys(mj_data) if CAPTURE_PARAMETERS == "all" else CAPTURE_PARAMETERS
+            
+        for key in keys:
+            try:
+                value = getattr(mj_data, key)
+                # Check if value is a numpy array and copy if needed
+                if isinstance(value, np.ndarray):
+                    self._d[key].append(value.copy())
+                elif np.isscalar(value):
+                    self._d[key].append(value)
+                elif hasattr(value, "copy") and callable(value.copy):
+                    # Copy if it has a copy method (e.g., MuJoCo's MjArray)
+                    self._d[key].append(value.copy())
+                else:
+                    self._d[key].append(value)
+            except AttributeError:
+                print_warning(f"Key '{key}' not found in MjData. Skipping.") if VERBOSITY else None
+            except Exception as e:
+                print(f"An error occurred while capturing '{key}': {e}") if VERBOSITY else None
+                
+    def unwrap(self) -> Dict[str, np.ndarray]:
+        """Unwrap the captured simulation data into a structured format with NumPy arrays."""
+        unwrapped_data = {}
+
+        for key, value_list in self._d.items():
+            if not value_list:  # Skip empty lists
+                unwrapped_data[key] = np.array([])
+                continue
+
+            try:
+                # Check if all items in the list have the same shape for array data
+                if isinstance(value_list[0], np.ndarray):
+                    # Check if arrays have consistent shapes
+                    shapes = [arr.shape for arr in value_list]
+                    if all(shape == shapes[0] for shape in shapes):
+                        unwrapped_data[key] = np.stack(value_list)
+                    else:
+                        # For arrays with different shapes, keep as a list
+                        unwrapped_data[key] = value_list
+                else:
+                    # Convert to a NumPy array if it's a list of scalars
+                    unwrapped_data[key] = np.array(value_list)
+            except ValueError as e:
+                print(f"ValueError while unwrapping key '{key}': {e}. Keeping as list.")
+                unwrapped_data[key] = value_list  # Store as a list if conversion fails
+            except Exception as e:
+                print(f"Error processing key '{key}': {e}")
+                unwrapped_data[key] = value_list
+
+        return unwrapped_data
+
+    @property
+    def shape(self):
+        """TODO: Implement a method to return the shape of the captured data."""
+        if not self._d:
+            return None
+
+        self.unwrap()  # Ensure data is unwrapped before checking shape
+
+        for key, value in self._d.items():
+            try:
+                if isinstance(value, np.ndarray):
+                    print(f"{key}: {value.shape}")
+                elif isinstance(value, list):
+                    print(f"{key}: {len(value)}")
+                elif isinstance(value, dict):
+                    print(f"{key}: {len(value)}")
+                else:  # Handle any other types
+                    print(f"{key}: {type(value)}")
+
+            except Exception:
+                pass
+
+    def __del__(self):
+        self._d.clear()
+        import gc
+
+        gc.collect()
+
+    def __len__(self):
+        if not self._d:
+            return 0
+        # Return the length of one of the data lists (assuming all have same length)
+        return len(next(iter(self._d.values())))
+
+    def __str__(self):
+        return f"{self.__class__.__name__}({len(self)} Step(s) Captured)"
+
+    def __repr__(self):
+        return self.__str__()
+
+    @staticmethod
+    def _get_public_keys(obj):
+        """Get all public keys of an object."""
+        return [
+            name
+            for name in dir(obj)
+            if not name.startswith("_") and not callable(getattr(obj, name))
+        ]

mujoco_toolbox/__init__.py

@@ -0,0 +1,67 @@
+"""
+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.
+
+Constants:
+----------
+- CAPTURE_PARAMETERS: List of MjData fields to capture during simulation.
+
+Attributes:
+-----------
+- __version__ (str): The current version of the package.
+- __author__ (str): The author of the package.
+- __all__ (list): List of public objects of the module, as interpreted by `from module import *`.
+
+Notes:
+------
+This package is still under development. Report any issues to https://github.com/MGross21/mujoco-toolbox/issues.
+"""
+
+from .Controller import (cosineController, randomController, sineController,
+                         stepController)
+from .Utils import print_warning
+from .Wrapper import Wrapper
+
+__version__ = "0.1.8"
+__author__ = "Michael Gross"
+__github_repo__ = "mujoco-toolbox"
+__license__ = "MIT"
+__status__ = "Development"
+
+# `from mujoco_toolbox import *` will import these objects
+__all__ = [
+    "Wrapper",
+    "sineController",
+    "cosineController",
+    "randomController",
+    "stepController",
+    "timer",
+]
+
+CAPTURE_PARAMETERS = [
+    "time",
+    "qpos",
+    "qvel",
+    "act",
+    "qacc",
+    "xpos",
+    "xquat",
+    "xmat",
+    "ctrl",
+    "sensordata",
+]  # MjData default fields to capture during simulation
+
+VERBOSITY = False
+
+if True:
+    print_warning(
+        f"{__package__} ({__version__}) is still under development.",
+        f"Report any issues to https://github.com/MGross21/{__github_repo__}/issues\n",
+    )

mujoco_toolbox-0.1.8.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Michael Gross
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mujoco_toolbox-0.1.8.dist-info/METADATA

@@ -0,0 +1,186 @@
+Metadata-Version: 2.3
+Name: mujoco-toolbox
+Version: 0.1.8
+Summary: A modern package to handle MuJoCo environments
+License: MIT
+Keywords: MuJoCo,robotics,simulation,toolbox,wrapper
+Author: MGross21
+Author-email: MGross21@TODO.com
+Requires-Python: >=3.10
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Topic :: Software Development :: Libraries
+Provides-Extra: dev
+Requires-Dist: ffmpeg
+Requires-Dist: ipywidgets
+Requires-Dist: matplotlib
+Requires-Dist: mediapy
+Requires-Dist: mujoco
+Requires-Dist: mypy ; extra == "dev"
+Requires-Dist: numpy
+Requires-Dist: pytest ; extra == "dev"
+Requires-Dist: pyyaml
+Requires-Dist: ruff ; extra == "dev"
+Requires-Dist: scipy
+Requires-Dist: screeninfo
+Requires-Dist: tqdm
+Requires-Dist: trimesh
+Project-URL: Documentation, https://github.com/MGross21/mujoco-toolbox
+Project-URL: Homepage, https://github.com/MGross21/mujoco-toolbox
+Project-URL: Repository, https://github.com/MGross21/mujoco-toolbox
+Project-URL: bug-tracker, https://github.com/MGross21/mujoco-toolbox/issues
+Description-Content-Type: text/markdown
+
+![Build](https://github.com/MGross21/mujoco-toolbox/actions/workflows/ci.yml/badge.svg)
+![Python](https://img.shields.io/badge/python-3.10%20|%203.11%20|%203.12%20|%203.13-blue)
+![License](https://img.shields.io/github/license/MGross21/mujoco-toolbox)
+[![PyPI](https://github.com/MGross21/mujoco-toolbox/actions/workflows/publish.yml/badge.svg)](https://github.com/MGross21/mujoco-toolbox/actions/workflows/publish.yml)
+
+# Mujoco Toolbox
+
+Streamlines the MuJoCo Physics Simulator
+
+## PyPI Package
+
+```bash
+pip install -U mujoco-toolbox
+```
+
+## For Local Install
+
+*Navigate to the Directory that you wish to place the package*
+
+<details>
+<summary><b>JUPYTER NOTEBOOK</b></summary>
+</br>
+*Place Commands Inside Cell*
+
+One-Time Install:
+
+```python
+!git clone https://github.com/MGross21/mujoco-toolbox
+!pip install -U ./mujoco-toolbox
+```
+
+Check for Updates:
+
+```python
+!git pull origin main
+```
+
+Install/Check for Updates:
+
+```python
+folder = "mujoco-toolbox"
+repo_url = f"https://github.com/MGross21/{folder}"
+
+!if [ -d "{folder}" ]; then cd {folder} && git pull && cd ..; else git clone {repo_url}; fi
+!pip install -U ./{folder}
+```
+
+</details>
+</br>
+
+<details>
+<summary><b>BASH WINDOW</b></summary>
+</br>
+One-Time Install:
+
+```bash
+git clone https://github.com/MGross21/mujoco-toolbox
+pip install -U ./mujoco-toolbox
+```
+
+Check for Updates:
+```bash
+git pull origin main
+```
+
+Install/Check for Updates:
+
+```bash
+REPO_DIR="mujoco-toolbox"
+REPO_URL="https://github.com/MGross21/mujoco-toolbox"
+
+if [ -d "$REPO_DIR" ]; then
+  cd "$REPO_DIR" && git fetch origin && git pull origin main
+else
+  git clone $REPO_URL
+  cd "$REPO_DIR"
+fi
+
+pip install -U ./mujoco-toolbox
+```
+
+</details>
+
+## Extra Packages
+
+<details>
+<summary><b>FFMPEG</b></summary>
+
+</br>
+
+*Required for [mediapy](https://google.github.io/mediapy/mediapy.html) dependency*
+
+**Windows**
+
+```bash
+winget install ffmpeg
+ffmpeg -version
+```
+
+**Linux**
+
+```bash
+sudo apt update && sudo apt install ffmpeg
+ffmpeg -version
+```
+
+**MacOS**
+
+*Using Homebrew*
+
+```bash
+brew install ffmpeg
+ffmpeg -version
+```
+
+*Using MacPorts*
+
+```bash
+sudo port install ffmpeg
+ffmpeg -version
+```
+
+</details>
+
+## To Run in Scripts
+
+```python
+import mujoco_toolbox as mjtb
+from mujoco_toolbox import Wrapper
+
+# Optional Global Variables
+mjtb.VERBOSITY = True # Additional Performance and Debugging Print Statements (Default=False)
+mjtb.CAPTURE_PARAMETERS = "all" # MjData Simulation Parameters (Default = ['time', 'qpos', 'qvel', 'act', 'qacc', 'xpos', 'xquat', 'xmat', 'ctrl', 'sensordata']) 
+
+# Example Code
+Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
+```
+
+## Pre-Made Controllers
+
+```python
+import mujoco_toolbox as mjtb
+
+mjtb.sineController(m,d,**kwargs)
+mjtb.cosineController(m,d,**kwargs)
+mjtb.stepController(m,d,**kwargs)
+mjtb.randomController(m,d,**kwargs)
+
+# Wrapper can use custom controllers as well!
+```
+

mujoco_toolbox-0.1.8.dist-info/RECORD

@@ -0,0 +1,8 @@
+mujoco_toolbox/__init__.py,sha256=5F4JYO_gNSHdvBeu3F8cIr5V_3SeCKmtoMFZOcNswAY,1818
+mujoco_toolbox/Controller.py,sha256=-q4qt43IWw0wz7TT4FnL6omLX-W7166rAdltTWSOflw,4659
+mujoco_toolbox/Utils.py,sha256=eDGJ973ItipwA-qgrWpKTEJke1-Xx4oc3asCXQp25kc,1900
+mujoco_toolbox/Wrapper.py,sha256=UmSq9vzaI96C2gBXcbshMNviS4y3_a7he168WRpQOwc,29050
+mujoco_toolbox-0.1.8.dist-info/LICENSE,sha256=Fh6x3ZPSwg6wLoA0jW1U0U4cM5tAFm2hi1JMypVc1V8,1089
+mujoco_toolbox-0.1.8.dist-info/METADATA,sha256=pyNDcZtroQGNrOr_-v8gtUryVg37LIy8HQgAs1yDsM8,4071
+mujoco_toolbox-0.1.8.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+mujoco_toolbox-0.1.8.dist-info/RECORD,,

mujoco_toolbox-0.1.8.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.1
+Root-Is-Purelib: true
+Tag: py3-none-any