mujoco-toolbox 0.2.2__tar.gz → 0.3.3__tar.gz

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/PKG-INFO

@@ -1,20 +1,36 @@
 Metadata-Version: 2.3
 Name: mujoco-toolbox
-Version: 0.2.2
+Version: 0.3.3
 Summary: A modern package to handle MuJoCo environments
-License: MIT
-Keywords: MuJoCo,robotics,simulation,toolbox,wrapper
+License: 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.
+Keywords: MuJoCo,mujoco,robotics,simulation,toolbox,wrapper
 Author: MGross21
 Author-email: MGross21@TODO.com
 Requires-Python: >=3.10
-Classifier: Development Status :: 3 - Alpha
+Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
-Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
-Classifier: Programming Language :: Python :: 3.11
-Classifier: Programming Language :: Python :: 3.12
-Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Software Development :: Libraries
 Requires-Dist: colorama
 Requires-Dist: ffmpeg
@@ -29,9 +45,10 @@ Requires-Dist: screeninfo
 Requires-Dist: tqdm
 Requires-Dist: trimesh
 Project-URL: Bug Tracker, https://github.com/MGross21/mujoco-toolbox/issues
-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: 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: changelog, https://github.com/MGross21/mujoco-toolbox/blob/main/CHANGELOG.md
 Description-Content-Type: text/markdown
 
 ![Build](https://github.com/MGross21/mujoco-toolbox/actions/workflows/ci.yml/badge.svg)
@@ -97,27 +114,26 @@ ffmpeg -version
 
 </details>
 
-## To Run in Scripts
+## Example Script
+
+*Bare minimum to run MuJoCo simulation and display result*
 
 ```python
 import mujoco_toolbox as mjtb
 
-# Optional Global Variables
-mjtb.CAPTURE_PARAMETERS = "all" # MjData Simulation Parameters (Default = ['time', 'qpos', 'qvel', 'act', 'qacc', 'xpos', 'xquat', 'xmat', 'ctrl', 'sensordata']) 
-
-# Example Code
-mjtb.Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
+mjtb.Wrapper("path/to/your/xml").run(render=True).save()
 ```
 
 ## Pre-Made Controllers
 
 ```python
-from mujoco_toolbox import \
-sineController,\
-cosineController,\
-stepController,\
-randomController,\
-realTimeController
+from mujoco_toolbox.controller import (
+    cosine_controller,
+    random_controller,
+    live_controller,
+    sine_controller,
+    step_controller,
+)
 
 # Wrapper can use custom controllers as well!
 ```
@@ -131,3 +147,4 @@ realTimeController
 ### URDF
 
 ![UR5](assets/images/ur5_render_no_gui.png)
+

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/README.md

@@ -61,27 +61,26 @@ ffmpeg -version
 
 </details>
 
-## To Run in Scripts
+## Example Script
+
+*Bare minimum to run MuJoCo simulation and display result*
 
 ```python
 import mujoco_toolbox as mjtb
 
-# Optional Global Variables
-mjtb.CAPTURE_PARAMETERS = "all" # MjData Simulation Parameters (Default = ['time', 'qpos', 'qvel', 'act', 'qacc', 'xpos', 'xquat', 'xmat', 'ctrl', 'sensordata']) 
-
-# Example Code
-mjtb.Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
+mjtb.Wrapper("path/to/your/xml").run(render=True).save()
 ```
 
 ## Pre-Made Controllers
 
 ```python
-from mujoco_toolbox import \
-sineController,\
-cosineController,\
-stepController,\
-randomController,\
-realTimeController
+from mujoco_toolbox.controller import (
+    cosine_controller,
+    random_controller,
+    live_controller,
+    sine_controller,
+    step_controller,
+)
 
 # Wrapper can use custom controllers as well!
 ```
@@ -94,4 +93,4 @@ realTimeController
 
 ### URDF
 
-![UR5](assets/images/ur5_render_no_gui.png)
+![UR5](assets/images/ur5_render_no_gui.png)

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/mujoco_toolbox/__init__.py

@@ -3,13 +3,14 @@
 A toolbox for working with MuJoCo simulations. This package provides various utilities and controllers to facilitate
 the simulation process.
 
-Modules:
+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.
+- `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:
+Constants
 ----------
 - CAPTURE_PARAMETERS: List of MjData fields to capture during simulation.
 
@@ -19,13 +20,19 @@ This package is still under development. Report any issues to https://github.com
 
 """
 
-from .assets import WORLD_ASSETS, GloveBox
+from .assets import WORLD_ASSETS, glovebox
 from .builder import Builder
-from .controller import cosineController, randomController, realTimeController, sineController, stepController
-from .utils import _Platform, print_success, print_warning, timer
+from .controller import (
+    cosine_controller,
+    live_controller,
+    random_controller,
+    sine_controller,
+    step_controller,
+)
+from .utils import _Platform
 from .wrapper import Wrapper
 
-__version__ = "0.2.2"
+__version__ = "0.3.3"
 __author__ = "Michael Gross"
 __github_repo__ = "mujoco-toolbox"
 __license__ = "MIT"
@@ -35,18 +42,18 @@ __all__ = [
     "COMPUTER",
     "WORLD_ASSETS",
     "Builder",
-    "GloveBox",
     "Wrapper",
-    "cosineController",
-    "randomController",
-    "realTimeController",
-    "sineController",
-    "stepController",
-    "timer",
+    "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",
@@ -60,8 +67,11 @@ CAPTURE_PARAMETERS = [  # MjData default fields to capture during simulation
     "sensordata",
 ]
 
-if __version__.split(".")[0] == "0":
-    print_warning(
+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-0.2.2 → mujoco_toolbox-0.3.3}/mujoco_toolbox/assets.py

@@ -10,7 +10,7 @@ WORLD_ASSETS = """
   </asset>
 """
 
-def GloveBox(width=1.25, depth=0.75, height=1.0, glass_thickness=0.05) -> str:
+def glovebox(width=1.25, depth=0.75, height=1.0, glass_thickness=0.05) -> str:
     return f"""
 <mujoco>
     <asset>

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/mujoco_toolbox/builder.py

@@ -3,11 +3,44 @@ import xml.etree.ElementTree as ET
 
 
 class Builder:
-    def __init__(self, xml_input: str) -> None:
+    """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."""
-        self.tree, self.root = self.load_model(xml_input)
+        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):
+    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"
@@ -22,7 +55,7 @@ class Builder:
 
                 # If it's a URDF (robot tag), wrap it in mujoco
                 if root.tag == "robot":
-                    return self.wrap_urdf_in_mujoco(tree)
+                    return self._wrap_urdf_in_mujoco(tree)
 
                 # If root is not mujoco, wrap it
                 if root.tag != "mujoco":
@@ -42,7 +75,7 @@ class Builder:
 
             # Handle URDF files
             if xml_input.lower().endswith(".urdf"):
-                return self.wrap_urdf_in_mujoco(ET.parse(xml_input))
+                return self._wrap_urdf_in_mujoco(ET.parse(xml_input))
 
             # Parse regular XML file
             try:
@@ -60,7 +93,7 @@ class Builder:
                 msg = f"Invalid XML file: {xml_input}"
                 raise ValueError(msg)
 
-    def wrap_urdf_in_mujoco(self, urdf_input):
+    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.
         """
@@ -101,7 +134,7 @@ class Builder:
         # 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:
+    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)
@@ -125,7 +158,7 @@ class Builder:
             model_root = other.root
         elif isinstance(other, str):
             # If the input is a string, load as XML
-            _, model_root = self.load_model(other)
+            _, model_root = self._load_model(other)
         else:
             msg = "Addition only supported between Builder objects or with XML strings/files"
             raise TypeError(msg)
@@ -136,7 +169,7 @@ class Builder:
 
         # Merge all relevant sections
         for section in sections_to_merge:
-            self.merge_tags(section, self.root, model_root)
+            self._merge_tags(section, self.root, model_root)
 
         # Return self to allow for method chaining
         return self
@@ -147,7 +180,7 @@ class Builder:
             return self
         return self.__add__(other)
 
-    def save(self, file_path) -> None:
+    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
@@ -156,6 +189,7 @@ class Builder:
         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."""
@@ -185,7 +219,11 @@ class Builder:
         """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):
-        """Property to quickly grab the XML string of the model."""
+        """XML string of the model."""
         return str(self)

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/mujoco_toolbox/controller.py

@@ -29,7 +29,7 @@ def _apply_control(model, data, value, joint=None, axis=None, delay=0) -> None:
         data.qpos[axis] = value
 
 
-def stepController(model, data, **kwargs) -> None:
+def step_controller(model, data, **kwargs) -> None:
     """A step controller for the simulation.
 
     Args:
@@ -57,7 +57,7 @@ def stepController(model, data, **kwargs) -> None:
     _apply_control(model, data, amplitude, joint=joint, axis=axis, delay=delay)
 
 
-def sineController(model, data, **kwargs) -> None:
+def sine_controller(model, data, **kwargs) -> None:
     """A simple sine wave controller for the simulation.
 
     Args:
@@ -85,7 +85,7 @@ def sineController(model, data, **kwargs) -> None:
     _apply_control(model, data, value, joint=joint, delay=delay)
 
 
-def cosineController(model, data, **kwargs) -> None:
+def cosine_controller(model, data, **kwargs) -> None:
     """A simple cosine wave controller for the simulation.
 
     Args:
@@ -113,7 +113,7 @@ def cosineController(model, data, **kwargs) -> None:
     _apply_control(model, data, value, joint=joint, delay=delay)
 
 
-def randomController(model, data, **kwargs) -> None:
+def random_controller(model, data, **kwargs) -> None:
     """A random controller for the simulation.
 
     Args:
@@ -141,7 +141,7 @@ def randomController(model, data, **kwargs) -> None:
     value = amplitude * random.rand()
     _apply_control(model, data, value, joint=joint, axis=axis, delay=delay)
 
-def realTimeController(model, data, **kwargs) -> None:
+def live_controller(model, data, **kwargs) -> None:
     """A real-time controller for the simulation.
 
     Args:
@@ -151,12 +151,12 @@ def realTimeController(model, data, **kwargs) -> None:
         None
 
     """
-    from .utils import print_warning
+    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...")
+            _print_warning(f"'{key}' is not a valid attribute of MjData. Skipping...")
 
     # untested
     # if hasattr(data, 'control'):

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/mujoco_toolbox/utils.py

@@ -20,7 +20,7 @@ def timer(func):
 
     return time_wrapper
 
-def print_warning(*args, prefix: bool = True) -> None:
+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:
@@ -49,7 +49,7 @@ def print_warning(*args, prefix: bool = True) -> None:
             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:
+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:
@@ -94,12 +94,12 @@ class _Platform:
     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"CPU Count: {self.CPU_COUNT}\n"
-            f"Architecture: {self.ARCHITECTURE}\n"
             f"Python Implementation: {self.PYTHON_IMPLEMENTATION}\n"
             f"Resolution (disp0): {self.RESOLUTION}\n"
             f"GUI Enabled: {self.GUI_ENABLED}"
@@ -121,5 +121,14 @@ class _Platform:
             monitor = get_monitors()[0]
             return (monitor.width, monitor.height), True
         except Exception:
-            print_warning("Detected headless operation. Disabling GUI...")
+            _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

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/mujoco_toolbox/wrapper.py

@@ -4,10 +4,10 @@ import threading
 import time
 import xml.etree.ElementTree as ET
 from collections import defaultdict
+from collections.abc import Callable
 from functools import lru_cache
-from typing import Any, Callable, Optional, TypeAlias
+from typing import Any, TypeAlias
 
-import matplotlib.pyplot as plt
 import mediapy as media
 import mujoco
 import mujoco.viewer
@@ -15,7 +15,7 @@ import numpy as np
 import trimesh
 import yaml
 
-from .utils import print_warning
+from .utils import _print_warning
 
 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."
@@ -27,17 +27,23 @@ mjData: TypeAlias = mujoco.MjData
 class Wrapper:
     """A class to handle MuJoCo simulations."""
 
-    def __init__(self, xml:str, duration:int=10, fps:int=30, resolution:tuple[int,int]=(400,300), initialConditions:Optional[dict[str, list]]=None, controller:Optional[Callable[[mjModel, mjData, Any], None]]=None, *args, **kwargs) -> None:
+    def __init__(self,
+                 xml:str,
+                 duration:int=10,
+                 data_rate:int=100,
+                 fps:int=30,
+                 resolution:tuple[int,int]=(400,300),
+                 initial_conditions:dict[str, list] | None=None,
+                 controller:Callable[[mjModel, mjData, Any], None] | None=None,
+                 *args,
+                 **kwargs) -> None:
         # xml = "<mujoco></mujoco>" if xml.strip() == "<mujoco/>" else xml
-        if initialConditions is None:
-            initialConditions = {}
-        assert xml is not None, "XML file or string is required to initialize the Wrapper."
         self._load_model(xml, **kwargs)
 
         self.duration = duration
         self.fps = fps
+        self.data_rate = data_rate
         self.resolution = resolution  # recursively sets width and height
-        self.init_conditions = initialConditions
         self.controller = controller
 
         # Predefined simulation parameters but can be overridden
@@ -46,6 +52,7 @@ class Wrapper:
         self.gravity = kwargs.get("gravity", self._model.opt.gravity)
 
         self._data = mujoco.MjData(self._model)
+        self.init_conditions = initial_conditions or {} # MUST BE AFTER DATA INITIALIZATION
 
         # Auto-Populate the names of bodies, joints, and actuators
         self._body_names = [self._model.body(i).name for i in range(self._model.nbody)]
@@ -58,11 +65,10 @@ class Wrapper:
     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):
+            if os.path.exists(xml):
+                # Convert the XML path to an absolute path
+                xml_path = os.path.abspath(xml)
                 # Extract and validate file extension
                 extension = os.path.splitext(xml_path)[1].lower()[1:]
 
@@ -72,33 +78,24 @@ class Wrapper:
                     self._load_urdf_file(xml_path, **kwargs)
                 else:
                     msg = f"Unsupported file extension: '{extension}'. Please provide an XML or URDF file."
-                    raise ValueError(
-                        msg,
-                    )
-            # If the file doesn't exist, assume it's a string and attempt to load it as XML
-            elif "<mujoco>" in xml or "<mujoco/>" in xml or "<robot>" in xml:
-                self._load_xml_string(xml)
+                    raise ValueError(msg)
             else:
-                msg = f"Model file not found: {xml_path}. Ensure the file path is correct and accessible."
-                raise FileNotFoundError(
-                    msg,
-                )
-
+                # If the file doesn't exist, assume it's a string and attempt to load it as XML
+                try:
+                    ET.fromstring(xml)  # Try parsing the string as XML
+                    self._load_xml_string(xml)
+                except ET.ParseError:
+                    msg = f"Model not able to be loaded from the provided XML string: {xml}"
+                    raise ValueError(msg)
         except FileNotFoundError as e:
             msg = f"Failed to load the MuJoCo model: {e}"
             raise FileNotFoundError(msg) from e
-
         except ValueError as e:
             msg = f"Invalid value encountered while loading the model: {e}"
-            raise ValueError(
-                msg,
-            ) from e
-
+            raise ValueError(msg) from e
         except Exception as e:
             msg = f"Unexpected error while loading the MuJoCo model: {e}"
-            raise Exception(
-                msg,
-            )
+            raise Exception(msg)
 
     def _load_xml_file(self, xml: str, **kwargs) -> None:
         """Load a MuJoCo model from an XML file."""
@@ -119,6 +116,7 @@ class Wrapper:
 
     def _load_urdf_file(self, urdf_path: str, **kwargs: Any) -> None:
         """Process and load a URDF file for use with MuJoCo."""
+        # BUG: Generate actuators from joints in URDF so the controller can actually work
 
         def convert_dae_to_stl(meshdir: str) -> None:
             """Convert all DAE files in a directory (including subdirectories) to STL."""
@@ -241,13 +239,13 @@ class Wrapper:
                 thread.join()
 
     @property
-    def model(self) -> mujoco.MjModel:
+    def model(self) -> 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."""
+    def data(self) -> mjData:
+        """Read-only property to access the MjData single-step object. Use `captured_data` to access the entire simulation data."""
         return self._data
 
     @property
@@ -272,13 +270,14 @@ class Wrapper:
 
     @frames.deleter
     def frames(self) -> None:
+        # BUG: This is not working as intended
         self._frames.clear()
         import gc
-
         gc.collect()
 
     @property
     def duration(self) -> float:
+        """Duration of the simulation in seconds."""
         return self._duration
 
     @duration.setter
@@ -290,6 +289,7 @@ class Wrapper:
 
     @property
     def fps(self) -> float:
+        """Frames per second."""
         return self._fps
 
     @fps.setter
@@ -301,6 +301,7 @@ class Wrapper:
 
     @property
     def resolution(self) -> tuple[int, int]:
+        """Resolution of the simulation in pixels."""
         return self._resolution
 
     @resolution.setter
@@ -326,27 +327,33 @@ class Wrapper:
         self._model.vis.Global.offheight = self._height
 
     @property
-    def init_conditions(self):
+    def initial_conditions(self):
+        """Initial conditions for the simulation."""
         return self._initcond
 
-    @init_conditions.setter
-    def init_conditions(self, values) -> None:
+    @initial_conditions.setter
+    def initial_conditions(self, values) -> None:
         if not isinstance(values, dict):
             msg = "Initial conditions must be a dictionary."
             raise ValueError(msg)
-        invalid_keys = [
-            key for key in values if not hasattr(mujoco.MjData(self._model), key)
-        ]
+
+        invalid_keys = [key for key in values if not hasattr(mujoco.MjData(self._model), key)]
         if invalid_keys:
-            SimulationData._get_public_keys(self._data)
+            _SimulationData._get_public_keys(self._data)
             msg = f"Invalid initial condition attributes: {', '.join(invalid_keys)}"
-            raise ValueError(
-                msg,
-            )
+            raise ValueError(msg)
+
         self._initcond = values
 
+        # Apply initial conditions directly when set
+        for key, value in values.items():
+            if hasattr(self._data, key):
+                setattr(self._data, key, value)
+            else:
+                _print_warning(f"'{key}' is not a valid attribute of MjData.")
+
     @property
-    def controller(self) -> Optional[Callable[[mjModel, mjData, Any], None]]:
+    def controller(self) -> Callable[[mjModel, mjData, Any], None] | None:
         """Controller Function."""
         return self._controller
 
@@ -359,6 +366,7 @@ class Wrapper:
 
     @property
     def ts(self) -> float:
+        """Timestep of the simulation in seconds."""
         return self._model.opt.timestep
 
     @ts.setter
@@ -370,17 +378,18 @@ class Wrapper:
 
     @property
     def data_rate(self) -> int:
+        """Data rate of the simulation in frames per second."""
         try:
             return self._dr
         except AttributeError:
-            print_warning(f"Use '{self.runSim.__name__}' first in order to access this value.")
+            _print_warning(f"Use '{self.run.__name__}' first in order to access this value.")
             return None
 
     @data_rate.setter
-    def data_rate(self, value) -> None:
+    def data_rate(self, value: int) -> None:
         if isinstance(value, float) and not isinstance(value, int):
             value = round(value)
-            print_warning(f"Data rate must be an integer. Rounding to the nearest integer ({value}).")
+            _print_warning(f"Data rate must be an integer. Rounding to the nearest integer ({value}).")
         if value <= 0:
             msg = "Data rate must be greater than 0."
             raise ValueError(msg)
@@ -393,23 +402,21 @@ class Wrapper:
 
     @property
     def gravity(self):
+        """Gravity vector of the simulation."""
         return self._model.opt.gravity
 
     @gravity.setter
-    def gravity(self, values) -> None:
-        if len(values) != 3:
+    def gravity(self, values: list | tuple | np.ndarray) -> None:
+        if not isinstance(values, (list, tuple, np.ndarray)) or len(values) != 3:
             msg = "Gravity must be a 3D vector."
             raise ValueError(msg)
-        self._model.opt.gravity = values
-
-    def _setInitialConditions(self) -> None:
-        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.")
+        self._model.opt.gravity = np.array(values)
 
-    def runSim(self, render: bool = False, camera: Optional[str] = None, data_rate: int = 100, interactive: bool = False, multi_thread: bool = False) -> "Wrapper":
+    def run(self,
+            render: bool = False,
+            camera: str | None = None,
+            interactive: bool = False,
+            multi_thread: bool = False) -> "Wrapper":
         """Run the simulation with optional rendering and controlled data capture.
 
         Args:
@@ -428,26 +435,19 @@ class Wrapper:
         try:
             mujoco.set_mjcb_control(self._controller) if self._controller else None
             mujoco.mj_resetData(self._model, self._data)
-            self._setInitialConditions()
 
             # self._frames = [None] * (self.duration * self.fps + 1)
             frames = []
-            sim_data = SimulationData()
+            sim_data = _SimulationData()
             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
-            h = self._height
-            w = self._width
             dur = self._duration
-            num_geoms = self._geom_names.__len__()
 
-            # TODO: Fix references to data rate in future.
-            # May be desired to specify in object creation
-            self._dr = data_rate
-            capture_rate = data_rate * self.ts
+            capture_rate = self.data_rate * self.ts
             capture_interval = max(1, int(1.0 / capture_rate))
             render_interval = max(1, int(1.0 / (self._fps * self.ts)))
 
@@ -459,19 +459,22 @@ class Wrapper:
                 # TODO: Implement multi-threading
                 pass
 
-            from . import COMPUTER, MAX_GEOM_SCALAR
-            if COMPUTER.IDE == "jupyter":
-                from tqdm.notebook import tqdm as bar
+            from . import COMPUTER, MAX_GEOM_SCALAR, PROGRESS_BAR
+            from .utils import _EmptyContextManager
+
+            if PROGRESS_BAR and render:
+                if COMPUTER.IDE == "jupyter":
+                    from tqdm.notebook import tqdm as ProgressBar
+                else:
+                    from tqdm import tqdm as ProgressBar
             else:
-                from tqdm import tqdm as bar
-
-            with (
-                bar(total=total_steps,
-                      desc="Simulation",
-                      unit=" step", leave=False,
-                ) as pbar,
-                mujoco.Renderer(m, h, w, num_geoms * MAX_GEOM_SCALAR) as renderer,
-            ):
+                ProgressBar = _EmptyContextManager
+
+            # PEP Convention - Prevent Naming Conflict
+            _ProgressBar = ProgressBar(total=total_steps, desc="Simulation", unit=" step", leave=False)
+            _Renderer = mujoco.Renderer(self.model, self._height, self._width, self._geom_names.__len__() * MAX_GEOM_SCALAR)
+
+            with _ProgressBar as pbar, _Renderer as renderer:
                 step = 0
                 while d.time < dur:
                     mj_step(m, d)
@@ -489,7 +492,7 @@ class Wrapper:
                         # frame_count += 1
                         frames.append(renderer.render().copy())
 
-                    pbar.update(1)
+                        pbar.update(1) if PROGRESS_BAR else None
                     step += 1
 
                     # if verbose:
@@ -511,143 +514,196 @@ class Wrapper:
 
         return self
 
-    def _window(self) -> None:
+    def _window(self, show_menu: bool =True) -> None:
         """Open a window to display the simulation in real time."""
+        try:
+            m = self._model
+            d = self._data
+
+            def key_callback(key) -> bool:
+                return key in (27, ord("q"))  # 27 = ESC key, 'q' to quit
+
+            # NOTE: launch_passive may blocking despite docstring saying otherwise
+            _Viewer = mujoco.viewer.launch_passive(m, d,
+                                                    show_left_ui=show_menu,
+                                                    show_right_ui=show_menu,
+                                                    key_callback=key_callback)
+            with _Viewer as viewer:
+                viewer.sync()
+                start_time = time.time()
+
+                try:
+                    while viewer.is_running():
+                        current_time = time.time()
+                        dt = current_time - start_time  # Time difference between frames
+
+                        mujoco.mj_step(m, d)  # Advance simulation by one step
+                        viewer.sync()  # Sync the viewer
 
+                        start_time = current_time  # Reset start_time for the next frame
+                        time.sleep(max(0, 1.0 / self.fps - dt))  # Adjust sleep to match self.fps
 
+                except KeyboardInterrupt:
+                    viewer.close()
+        except Exception as e:
+            msg = "An error occurred while running the simulation."
+            raise RuntimeError(msg) from e
+        finally:
+            mujoco.set_mjcb_control(None)
 
     def liveView(self, show_menu: bool = True) -> None:
         """Open a window to display the simulation in real time."""
-        # BUG: Generate actuators from joints in URDF so the controller can actually work
+        # Run the window in a separate thread
+        gui = threading.Thread(target=self._window, kwargs={"show_menu": show_menu})
+        gui.start()
 
-        def window() -> None:
-            try:
-                m = self._model
-                d = self._data
-
-                def key_callback(key) -> bool:
-                    if key in (27, ord("q")):  # 27 = ESC key, 'q' to quit
-                        return True
-                    return False
-
-                # NOTE: launch_passive is blocking despite docstring saying otherwise
-                with mujoco.viewer.launch_passive(m, d,
-                                                show_left_ui=show_menu,
-                                                show_right_ui=show_menu,
-                                                key_callback=key_callback) as viewer:
-                    viewer.sync()
-                    start_time = time.time()
-
-                    try:
-                        while viewer.is_running():
-                            current_time = time.time()
-                            dt = current_time - start_time  # Time difference between frames
-
-                            mujoco.mj_step(m, d)  # Advance simulation by one step
-                            viewer.sync()  # Sync the viewer
-
-                            start_time = current_time  # Reset start_time for the next frame
-                            # Sleep to match real-time simulation speed
-                            time.sleep(max(0, 0.01 - dt))  # Adjust sleep to match real-time
-                    except KeyboardInterrupt:
-                        viewer.close()
-            except Exception as e:
-                msg = "An error occurred while running the simulation."
-                raise RuntimeError(msg) from e
-            finally:
-                mujoco.set_mjcb_control(None)
-
-        viewer = threading.Thread(target=window)
-        viewer.start()
-
-    def renderFrame(self, t=0, frame=0, title=None) -> Optional[str]:
-        """Render a specific frame as an image.
+    def _validate_and_extract_frames(
+        self,
+        frame_idx: int | tuple[int, int] | None = None,
+        time_idx: float | tuple[float, float] | None = None,
+    ) -> list[Any]:
+        """Validate and extract frames based on frame or time indices.
 
         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.
+            frame_idx (int or tuple, optional): Single frame index or (start, stop) frame indices.
+            time_idx (float or tuple, optional): Single time or (start, end) times in seconds.
 
         Returns:
-            None
+            List of frames
+
+        Raises:
+            ValueError: For invalid input parameters.
 
         """
         if not self._frames:
             msg = "No frames captured to render."
             raise ValueError(msg)
 
-        if t < 0 or frame < 0:
-            msg = "Time and frame index must be greater than or equal to zero."
-            raise ValueError(
-                msg,
-            )
-
-        if frame and t:
-            msg = "Can only specify singular time or frame parameter"
+        # Validate input parameters
+        if frame_idx is not None and time_idx is not None:
+            msg = "Can only specify either frame_idx or time_idx, not both."
             raise ValueError(msg)
 
-        try:
-            if t > 0:
-                frame = self.t2f(t)  # Convert time to frame index
+        # If both are None, use all frames
+        if frame_idx is None and time_idx is None:
+            return self._frames
+
+        # Handle time index conversion
+        if time_idx is not None:
+            if isinstance(time_idx, (int, float)):
+                frame_idx = self.t2f(time_idx)
+            elif isinstance(time_idx, tuple):
+                frame_idx = (self.t2f(time_idx[0]), self.t2f(time_idx[1]))
             else:
-                frame = int(frame)
+                msg = "time_idx must be a number or a tuple of numbers."
+                raise ValueError(msg)
 
-            plt.imshow(self._frames[frame])
-            plt.axis("off")
-            plt.title(title or f"Frame {frame}", loc="center")
-            plt.show()
+        # Convert single index to tuple range
+        if isinstance(frame_idx, (int, float)):
+            frame_idx = (frame_idx, frame_idx + 1)
 
-        except IndexError as e:
-            msg = f"Invalid frame index: {e}"
+        # Validate frame indices
+        start, stop = frame_idx
+        max_frames = len(self._frames)
+
+        if start < 0:
+            msg = f"Start index must be non-negative. Got {start}."
             raise ValueError(msg)
-        except TypeError as e:
-            msg = f"Invalid type for time or frame: {e}"
+        if stop > max_frames:
+            msg = f"Stop index must not exceed total frames ({max_frames}). Got {stop}."
             raise ValueError(msg)
-        except Exception as e:
-            msg = f"Unexpected error while rendering frame: {e}"
-            raise RuntimeError(msg)
+        if start >= stop:
+            msg = f"Start index ({start}) must be less than stop index ({stop})."
+            raise ValueError(msg)
+
+        # Select subset of frames
+        return self._frames[start:stop]
 
-    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.
+    def show(
+        self,
+        title: str | None = None,
+        codec: str = "gif",
+        frame_idx: int | tuple[int, int] | None = None,
+        time_idx: float | tuple[float, float] | None = None,
+    ) -> None:
+        """Render specific frame(s) as a video or GIF in a window.
 
         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.
+            title (str, optional): Title for the rendered media.
+            codec (str, optional): Video codec/format. Defaults to "gif".
+            frame_idx (int or tuple, optional): Single frame index or (start, stop) frame indices.
+            time_idx (float or tuple, optional): Single time or (start, end) times in seconds.
+
+        Raises:
+            ValueError: If no frames are captured or invalid input parameters.
 
         """
-        if not self._frames:
-            msg = "No frames captured to create media."
-            raise ValueError(msg)
-        # Display the media in a window
-        if not save:
+        try:
+            # Extract frames
+            subset_frames = self._validate_and_extract_frames(
+                frame_idx=frame_idx,
+                time_idx=time_idx,
+            )
+
+            # Show the video
             media.show_video(
-                self._frames,
-                fps=self._fps,
+                subset_frames,
+                fps=1 if len(subset_frames) == 1 else self._fps,
                 width=self._width,
                 height=self._height,
                 codec=codec,
                 title=title,
             )
-            return None
+        except Exception as e:
+            msg = "Error while showing video subset."
+            raise Exception(msg) from e
+
+    def save(
+        self,
+        title: str = "render",
+        codec: str = "gif",
+        frame_idx: int | tuple[int, int] | None = None,
+        time_idx: float | tuple[float, float] | None = None,
+    ) -> str:
+        """Save specific frame(s) as a video or GIF to a file.
+
+        Args:
+            title (str, optional): Filename for the saved media.
+            codec (str, optional): Video codec/format. Defaults to "gif".
+            frame_idx (int or tuple, optional): Single frame index or (start, stop) frame indices.
+            time_idx (float or tuple, optional): Single time or (start, end) times in seconds.
+
+        Returns:
+            str: Absolute path to the saved file.
+
+        Raises:
+            ValueError: If no frames are captured or invalid input parameters.
 
-        if not title.endswith(f".{codec}"):
-            title += f".{codec}"
+        """
+        try:
+            # Extract frames
+            subset_frames = self._validate_and_extract_frames(
+                frame_idx=frame_idx,
+                time_idx=time_idx,
+            )
 
-        # Save the frames to the specified file
-        available_codecs = ["gif", "mp4", "h264", "hevc", "vp9"]
-        if codec in available_codecs:
+            # Ensure the title ends with the correct codec extension
+            if not title.endswith(f".{codec}"):
+                title += f".{codec}"
+
+            # Save the video
             media.write_video(
-                title if not None else "render",
-                self._frames,
-                fps=self._fps,
+                title,
+                subset_frames,
+                fps=1 if len(subset_frames) == 1 else self._fps,
                 codec=codec,
             )
-        else:
-            msg = f"Unsupported codec '{codec}'. Supported codecs are {', '.join(available_codecs)}"
-            raise ValueError(msg)
 
-        return os.path.abspath(title)
+            return os.path.abspath(title)
+
+        except Exception as e:
+            msg = "Error while saving video subset."
+            raise Exception(msg) from e
 
     @lru_cache(maxsize=100)
     def t2f(self, t: float) -> int:
@@ -661,7 +717,7 @@ class Wrapper:
         """Convert frame index to time."""
         return frame / self._fps
 
-    def getBodyData(self, body_name: str, data_name: Optional[str] = None) -> np.ndarray:
+    def getBodyData(self, body_name: str, data_name: str | None = None) -> np.ndarray:
         """Get the data for a specific body in the simulation.
 
         Args:
@@ -727,7 +783,7 @@ class Wrapper:
             raise ValueError(msg) from e
 
 
-class SimulationData:
+class _SimulationData:
     """A class to store and manage simulation data."""
 
     __slots__ = ["_d"]

{mujoco_toolbox-0.2.2 → mujoco_toolbox-0.3.3}/pyproject.toml

@@ -1,48 +1,53 @@
 [build-system]
-requires = ["poetry-core"]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
 build-backend = "poetry.core.masonry.api"
 
-[tool.poetry]
+[project]
 name = "mujoco-toolbox"
-version = "0.2.2"
+version = "0.3.3"
 description = "A modern package to handle MuJoCo environments"
 readme = "README.md"
-license = "MIT"
+license = { file = "LICENSE" }
 packages = [{ include = "mujoco_toolbox" }]
-authors = ["MGross21 <MGross21@TODO.com>"]
-keywords = ["MuJoCo", "robotics", "simulation", "toolbox", "wrapper"]
+authors = [{ name = "MGross21", email = "MGross21@TODO.com" }]
+keywords = ["MuJoCo", "mujoco", "robotics", "simulation", "toolbox", "wrapper"]
 classifiers = [
-    "Development Status :: 3 - Alpha",
+    "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",
     "License :: OSI Approved :: MIT License",
     "Programming Language :: Python :: 3.10",
     "Topic :: Software Development :: Libraries"
 ]
+requires-python = ">=3.10"
+dependencies = [
+    "mujoco",
+    "numpy",
+    "scipy",
+    "matplotlib",
+    "pyyaml",
+    "mediapy",
+    "tqdm",
+    "ipywidgets",
+    "ffmpeg",
+    "trimesh",
+    "screeninfo",
+    "colorama"
+]
+
+# [project.extras]
+# beautified = ["colorama","ipywidgets","tqdm"]
+# gym = ["gym"]
 
-[tool.poetry.urls]
+[project.urls]
 homepage = "https://github.com/MGross21/mujoco-toolbox"
 repository = "https://github.com/MGross21/mujoco-toolbox"
 documentation = "https://github.com/MGross21/mujoco-toolbox"
+changelog = "https://github.com/MGross21/mujoco-toolbox/blob/main/CHANGELOG.md"
 "Bug Tracker" = "https://github.com/MGross21/mujoco-toolbox/issues"
 
-[tool.poetry.dependencies]
-python = ">=3.10"
-mujoco = "*"
-numpy = "*"
-scipy = "*"
-matplotlib = "*"
-pyyaml = "*"
-mediapy = "*"
-tqdm = "*"
-ipywidgets = "*"
-ffmpeg = "*"
-trimesh = "*"
-screeninfo = "*"
-colorama = "*"
-
 [tool.poetry.group.dev.dependencies]
 pytest = "*"
-ruff = "^0.9.9"
+ruff = "^0.11.2"
 mypy = "^1.15.0"
 pytypes = "^1.0b10"
 black = "^25.1.0"
@@ -54,14 +59,14 @@ pylint = "^3.3.4"
 [tool.semantic_release]
 version_variables = [
     "mujoco_toolbox/__init__.py:__version__",
-    "pyproject.toml:tool.poetry.version",
+    "pyproject.toml:project.version",
     "CITATION.cff:version"
 ]
-branch = "main"
+branch = "main"                 # Branch to track for triggering the release process
 changelog_file = "CHANGELOG.md"
 build_command = "pip install poetry && poetry build"
-dist_path = "dist/"
-create_pr = true
+dist_path = "dist/"             # Directory where the distribution files will be stored
+create_pr = true                # Option to create a pull request when a new release version is created
 pr_branch_prefix = "release/"
 upload_to_pypi = true
 upload_to_release = true
@@ -72,8 +77,12 @@ commit_message = "chore(release): {version} [skip ci]"
 commit_author = "semantic-release <semantic-release@example.com>"
 
 [tool.semantic_release.plugins]
-commit_analyzer = "@semantic-release/commit-analyzer"
-release_notes_generator = "@semantic-release/release-notes-generator"
+commit_analyzer = "@semantic-release/commit-analyzer"           # Analyzes commit messages to determine version bumps
+release_notes_generator = "@semantic-release/release-notes-generator"  # Generates release notes based on commit messages
+changelog = "@semantic-release/changelog"                       # Handles changelog generation
+git = "@semantic-release/git"                                   # Commits changes (version bump, changelog) back to the repo
+github = "@semantic-release/github"                             # Creates GitHub releases (tags) and adds release notes
+pypi = "@semantic-release/pypi"                                 # Uploads the distribution to PyPI
 
 [tool.semantic_release.commit_analyzer]
 preset = "angular"
@@ -111,7 +120,7 @@ preset_config = {types = [
     ]}
 
 [tool.ruff]
-line-length = 120
+line-length = 79 # PEP8 line length
 
 [tool.ruff.lint]
 # select = ["E", "F", "W", "C90"]  # Select error, warning, and flake8 rules