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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.3
2
2
  Name: mujoco-toolbox
3
- Version: 0.2.2
3
+ Version: 0.3.0
4
4
  Summary: A modern package to handle MuJoCo environments
5
5
  License: MIT
6
6
  Keywords: MuJoCo,robotics,simulation,toolbox,wrapper
@@ -10,11 +10,7 @@ Requires-Python: >=3.10
10
10
  Classifier: Development Status :: 3 - Alpha
11
11
  Classifier: Intended Audience :: Developers
12
12
  Classifier: License :: OSI Approved :: MIT License
13
- Classifier: Programming Language :: Python :: 3
14
13
  Classifier: Programming Language :: Python :: 3.10
15
- Classifier: Programming Language :: Python :: 3.11
16
- Classifier: Programming Language :: Python :: 3.12
17
- Classifier: Programming Language :: Python :: 3.13
18
14
  Classifier: Topic :: Software Development :: Libraries
19
15
  Requires-Dist: colorama
20
16
  Requires-Dist: ffmpeg
@@ -29,9 +25,9 @@ Requires-Dist: screeninfo
29
25
  Requires-Dist: tqdm
30
26
  Requires-Dist: trimesh
31
27
  Project-URL: Bug Tracker, https://github.com/MGross21/mujoco-toolbox/issues
32
- Project-URL: documentation, https://github.com/MGross21/mujoco-toolbox
33
- Project-URL: homepage, https://github.com/MGross21/mujoco-toolbox
34
- Project-URL: repository, https://github.com/MGross21/mujoco-toolbox
28
+ Project-URL: Documentation, https://github.com/MGross21/mujoco-toolbox
29
+ Project-URL: Homepage, https://github.com/MGross21/mujoco-toolbox
30
+ Project-URL: Repository, https://github.com/MGross21/mujoco-toolbox
35
31
  Description-Content-Type: text/markdown
36
32
 
37
33
  ![Build](https://github.com/MGross21/mujoco-toolbox/actions/workflows/ci.yml/badge.svg)
@@ -97,27 +93,26 @@ ffmpeg -version
97
93
 
98
94
  </details>
99
95
 
100
- ## To Run in Scripts
96
+ ## Example Script
97
+
98
+ *Bare minimum to run MuJoCo simulation and display result*
101
99
 
102
100
  ```python
103
101
  import mujoco_toolbox as mjtb
104
102
 
105
- # Optional Global Variables
106
- mjtb.CAPTURE_PARAMETERS = "all" # MjData Simulation Parameters (Default = ['time', 'qpos', 'qvel', 'act', 'qacc', 'xpos', 'xquat', 'xmat', 'ctrl', 'sensordata'])
107
-
108
- # Example Code
109
- mjtb.Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
103
+ mjtb.Wrapper("path/to/your/xml").run(render=True).save()
110
104
  ```
111
105
 
112
106
  ## Pre-Made Controllers
113
107
 
114
108
  ```python
115
- from mujoco_toolbox import \
116
- sineController,\
117
- cosineController,\
118
- stepController,\
119
- randomController,\
120
- realTimeController
109
+ from mujoco_toolbox.controller import (
110
+ cosine_controller,
111
+ random_controller,
112
+ live_controller,
113
+ sine_controller,
114
+ step_controller,
115
+ )
121
116
 
122
117
  # Wrapper can use custom controllers as well!
123
118
  ```
@@ -131,3 +126,4 @@ realTimeController
131
126
  ### URDF
132
127
 
133
128
  ![UR5](assets/images/ur5_render_no_gui.png)
129
+
@@ -61,27 +61,26 @@ ffmpeg -version
61
61
 
62
62
  </details>
63
63
 
64
- ## To Run in Scripts
64
+ ## Example Script
65
+
66
+ *Bare minimum to run MuJoCo simulation and display result*
65
67
 
66
68
  ```python
67
69
  import mujoco_toolbox as mjtb
68
70
 
69
- # Optional Global Variables
70
- mjtb.CAPTURE_PARAMETERS = "all" # MjData Simulation Parameters (Default = ['time', 'qpos', 'qvel', 'act', 'qacc', 'xpos', 'xquat', 'xmat', 'ctrl', 'sensordata'])
71
-
72
- # Example Code
73
- mjtb.Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
71
+ mjtb.Wrapper("path/to/your/xml").run(render=True).save()
74
72
  ```
75
73
 
76
74
  ## Pre-Made Controllers
77
75
 
78
76
  ```python
79
- from mujoco_toolbox import \
80
- sineController,\
81
- cosineController,\
82
- stepController,\
83
- randomController,\
84
- realTimeController
77
+ from mujoco_toolbox.controller import (
78
+ cosine_controller,
79
+ random_controller,
80
+ live_controller,
81
+ sine_controller,
82
+ step_controller,
83
+ )
85
84
 
86
85
  # Wrapper can use custom controllers as well!
87
86
  ```
@@ -94,4 +93,4 @@ realTimeController
94
93
 
95
94
  ### URDF
96
95
 
97
- ![UR5](assets/images/ur5_render_no_gui.png)
96
+ ![UR5](assets/images/ur5_render_no_gui.png)
@@ -5,9 +5,10 @@ the simulation process.
5
5
 
6
6
  Modules:
7
7
  --------
8
- - Wrapper: Contains the Wrapper class for interfacing with MuJoCo.
9
- - Controller: Includes sineController, cosineController, and randomController for controlling simulations.
10
- - Builder: Contains the Builder class for creating and merging MuJoCo models.
8
+ - `Wrapper`: Contains the Wrapper class for interfacing with MuJoCo.
9
+ - `Controller`: Includes sineController, cosineController, and randomController for controlling simulations.
10
+ - `Builder`: Contains the Builder class for creating and merging MuJoCo models.
11
+ - `assets`: Contains pre-defined assets for building MuJoCo models.
11
12
 
12
13
  Constants:
13
14
  ----------
@@ -19,13 +20,19 @@ This package is still under development. Report any issues to https://github.com
19
20
 
20
21
  """
21
22
 
22
- from .assets import WORLD_ASSETS, GloveBox
23
+ from .assets import WORLD_ASSETS, glovebox
23
24
  from .builder import Builder
24
- from .controller import cosineController, randomController, realTimeController, sineController, stepController
25
- from .utils import _Platform, print_success, print_warning, timer
25
+ from .controller import (
26
+ cosine_controller,
27
+ live_controller,
28
+ random_controller,
29
+ sine_controller,
30
+ step_controller,
31
+ )
32
+ from .utils import _Platform, print_warning
26
33
  from .wrapper import Wrapper
27
34
 
28
- __version__ = "0.2.2"
35
+ __version__ = "0.3.0"
29
36
  __author__ = "Michael Gross"
30
37
  __github_repo__ = "mujoco-toolbox"
31
38
  __license__ = "MIT"
@@ -35,14 +42,13 @@ __all__ = [
35
42
  "COMPUTER",
36
43
  "WORLD_ASSETS",
37
44
  "Builder",
38
- "GloveBox",
39
45
  "Wrapper",
40
- "cosineController",
41
- "randomController",
42
- "realTimeController",
43
- "sineController",
44
- "stepController",
45
- "timer",
46
+ "cosine_controller",
47
+ "glovebox",
48
+ "live_controller",
49
+ "random_controller",
50
+ "sine_controller",
51
+ "step_controller",
46
52
  ]
47
53
 
48
54
  COMPUTER = _Platform() # Singleton Instance
@@ -60,8 +66,9 @@ CAPTURE_PARAMETERS = [ # MjData default fields to capture during simulation
60
66
  "sensordata",
61
67
  ]
62
68
 
63
- if __version__.split(".")[0] == "0":
69
+ if __version__.startswith("0"):
64
70
  print_warning(
65
71
  f"{__package__} (v{__version__}) is still under development.",
66
72
  f"Report any issues to https://github.com/MGross21/{__github_repo__}/issues",
67
73
  )
74
+ del _Platform, print_warning
@@ -10,7 +10,7 @@ WORLD_ASSETS = """
10
10
  </asset>
11
11
  """
12
12
 
13
- def GloveBox(width=1.25, depth=0.75, height=1.0, glass_thickness=0.05) -> str:
13
+ def glovebox(width=1.25, depth=0.75, height=1.0, glass_thickness=0.05) -> str:
14
14
  return f"""
15
15
  <mujoco>
16
16
  <asset>
@@ -3,9 +3,24 @@ import xml.etree.ElementTree as ET
3
3
 
4
4
 
5
5
  class Builder:
6
- def __init__(self, xml_input: str) -> None:
6
+ def __init__(self, *args: str) -> None:
7
7
  """Initialize with an XML string or a file path. Input is required."""
8
- self.tree, self.root = self.load_model(xml_input)
8
+ for arg in args:
9
+ if not isinstance(arg, str):
10
+ msg = "Input must be an XML string or a file path"
11
+ raise TypeError(msg)
12
+ if len(args) == 0:
13
+ msg = "Input is required to initialize the Builder"
14
+ raise ValueError(msg)
15
+ if len(args) > 1:
16
+ # Create a new Builder instance for each additional argument and merge them
17
+ self.tree, self.root = self.load_model(args[0])
18
+ for additional_arg in args[1:]:
19
+ additional_builder = Builder(additional_arg)
20
+ self.__add__(additional_builder)
21
+ else:
22
+ xml_input = args[0]
23
+ self.tree, self.root = self.load_model(xml_input)
9
24
 
10
25
  def load_model(self, xml_input):
11
26
  """Load a MuJoCo model from a file or XML string."""
@@ -147,7 +162,7 @@ class Builder:
147
162
  return self
148
163
  return self.__add__(other)
149
164
 
150
- def save(self, file_path) -> None:
165
+ def save(self, file_path) -> "os.path":
151
166
  """Save the merged model to a file."""
152
167
  if self.tree is not None:
153
168
  # Format the XML with proper indentation before saving
@@ -156,6 +171,7 @@ class Builder:
156
171
  else:
157
172
  msg = "No model loaded. Cannot save."
158
173
  raise ValueError(msg)
174
+ return os.path.abspath(file_path)
159
175
 
160
176
  def _indent_xml(self, elem, level=0) -> None:
161
177
  """Add proper indentation to make the XML file more readable."""
@@ -185,6 +201,10 @@ class Builder:
185
201
  """Return the XML string of the model."""
186
202
  return self.__str__()
187
203
 
204
+ def __len__(self) -> int:
205
+ """Return the number of elements in the model."""
206
+ return len(self.root) if self.root is not None else 0
207
+
188
208
  @property
189
209
  def xml(self):
190
210
  """Property to quickly grab the XML string of the model."""
@@ -29,7 +29,7 @@ def _apply_control(model, data, value, joint=None, axis=None, delay=0) -> None:
29
29
  data.qpos[axis] = value
30
30
 
31
31
 
32
- def stepController(model, data, **kwargs) -> None:
32
+ def step_controller(model, data, **kwargs) -> None:
33
33
  """A step controller for the simulation.
34
34
 
35
35
  Args:
@@ -57,7 +57,7 @@ def stepController(model, data, **kwargs) -> None:
57
57
  _apply_control(model, data, amplitude, joint=joint, axis=axis, delay=delay)
58
58
 
59
59
 
60
- def sineController(model, data, **kwargs) -> None:
60
+ def sine_controller(model, data, **kwargs) -> None:
61
61
  """A simple sine wave controller for the simulation.
62
62
 
63
63
  Args:
@@ -85,7 +85,7 @@ def sineController(model, data, **kwargs) -> None:
85
85
  _apply_control(model, data, value, joint=joint, delay=delay)
86
86
 
87
87
 
88
- def cosineController(model, data, **kwargs) -> None:
88
+ def cosine_controller(model, data, **kwargs) -> None:
89
89
  """A simple cosine wave controller for the simulation.
90
90
 
91
91
  Args:
@@ -113,7 +113,7 @@ def cosineController(model, data, **kwargs) -> None:
113
113
  _apply_control(model, data, value, joint=joint, delay=delay)
114
114
 
115
115
 
116
- def randomController(model, data, **kwargs) -> None:
116
+ def random_controller(model, data, **kwargs) -> None:
117
117
  """A random controller for the simulation.
118
118
 
119
119
  Args:
@@ -141,7 +141,7 @@ def randomController(model, data, **kwargs) -> None:
141
141
  value = amplitude * random.rand()
142
142
  _apply_control(model, data, value, joint=joint, axis=axis, delay=delay)
143
143
 
144
- def realTimeController(model, data, **kwargs) -> None:
144
+ def live_controller(model, data, **kwargs) -> None:
145
145
  """A real-time controller for the simulation.
146
146
 
147
147
  Args:
@@ -94,12 +94,12 @@ class _Platform:
94
94
  def __repr__(self) -> str:
95
95
  return (
96
96
  f"Operating System: {self.OS}\n"
97
+ f"Architecture: {self.ARCHITECTURE}\n"
97
98
  f"Processor: {self.PROCESSOR}\n"
98
99
  f"Machine Type: {self.MACHINE}\n"
100
+ f"CPU Count: {self.CPU_COUNT}\n"
99
101
  f"Python Version: {self.PYTHON_VERSION}\n"
100
102
  f"IDE: {self.IDE}\n"
101
- f"CPU Count: {self.CPU_COUNT}\n"
102
- f"Architecture: {self.ARCHITECTURE}\n"
103
103
  f"Python Implementation: {self.PYTHON_IMPLEMENTATION}\n"
104
104
  f"Resolution (disp0): {self.RESOLUTION}\n"
105
105
  f"GUI Enabled: {self.GUI_ENABLED}"
@@ -4,10 +4,10 @@ import threading
4
4
  import time
5
5
  import xml.etree.ElementTree as ET
6
6
  from collections import defaultdict
7
+ from collections.abc import Callable
7
8
  from functools import lru_cache
8
- from typing import Any, Callable, Optional, TypeAlias
9
+ from typing import Any, TypeAlias
9
10
 
10
- import matplotlib.pyplot as plt
11
11
  import mediapy as media
12
12
  import mujoco
13
13
  import mujoco.viewer
@@ -27,17 +27,23 @@ mjData: TypeAlias = mujoco.MjData
27
27
  class Wrapper:
28
28
  """A class to handle MuJoCo simulations."""
29
29
 
30
- 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:
30
+ def __init__(self,
31
+ xml:str,
32
+ duration:int=10,
33
+ data_rate:int=100,
34
+ fps:int=30,
35
+ resolution:tuple[int,int]=(400,300),
36
+ initial_conditions:dict[str, list] | None=None,
37
+ controller:Callable[[mjModel, mjData, Any], None] | None=None,
38
+ *args,
39
+ **kwargs) -> None:
31
40
  # xml = "<mujoco></mujoco>" if xml.strip() == "<mujoco/>" else xml
32
- if initialConditions is None:
33
- initialConditions = {}
34
- assert xml is not None, "XML file or string is required to initialize the Wrapper."
35
41
  self._load_model(xml, **kwargs)
36
42
 
37
43
  self.duration = duration
38
44
  self.fps = fps
45
+ self.data_rate = data_rate
39
46
  self.resolution = resolution # recursively sets width and height
40
- self.init_conditions = initialConditions
41
47
  self.controller = controller
42
48
 
43
49
  # Predefined simulation parameters but can be overridden
@@ -46,6 +52,7 @@ class Wrapper:
46
52
  self.gravity = kwargs.get("gravity", self._model.opt.gravity)
47
53
 
48
54
  self._data = mujoco.MjData(self._model)
55
+ self.init_conditions = initial_conditions or {} # MUST BE AFTER DATA INITIALIZATION
49
56
 
50
57
  # Auto-Populate the names of bodies, joints, and actuators
51
58
  self._body_names = [self._model.body(i).name for i in range(self._model.nbody)]
@@ -58,11 +65,10 @@ class Wrapper:
58
65
  def _load_model(self, xml: str, **kwargs: Any) -> None:
59
66
  """Load a MuJoCo model from an XML file or a string."""
60
67
  try:
61
- # Convert the XML path to an absolute path
62
- xml_path = os.path.abspath(xml)
63
-
64
68
  # Check if the path exists
65
- if os.path.exists(xml_path):
69
+ if os.path.exists(xml):
70
+ # Convert the XML path to an absolute path
71
+ xml_path = os.path.abspath(xml)
66
72
  # Extract and validate file extension
67
73
  extension = os.path.splitext(xml_path)[1].lower()[1:]
68
74
 
@@ -72,33 +78,24 @@ class Wrapper:
72
78
  self._load_urdf_file(xml_path, **kwargs)
73
79
  else:
74
80
  msg = f"Unsupported file extension: '{extension}'. Please provide an XML or URDF file."
75
- raise ValueError(
76
- msg,
77
- )
78
- # If the file doesn't exist, assume it's a string and attempt to load it as XML
79
- elif "<mujoco>" in xml or "<mujoco/>" in xml or "<robot>" in xml:
80
- self._load_xml_string(xml)
81
+ raise ValueError(msg)
81
82
  else:
82
- msg = f"Model file not found: {xml_path}. Ensure the file path is correct and accessible."
83
- raise FileNotFoundError(
84
- msg,
85
- )
86
-
83
+ # If the file doesn't exist, assume it's a string and attempt to load it as XML
84
+ try:
85
+ ET.fromstring(xml) # Try parsing the string as XML
86
+ self._load_xml_string(xml)
87
+ except ET.ParseError:
88
+ msg = f"Model not able to be loaded from the provided XML string: {xml}"
89
+ raise ValueError(msg)
87
90
  except FileNotFoundError as e:
88
91
  msg = f"Failed to load the MuJoCo model: {e}"
89
92
  raise FileNotFoundError(msg) from e
90
-
91
93
  except ValueError as e:
92
94
  msg = f"Invalid value encountered while loading the model: {e}"
93
- raise ValueError(
94
- msg,
95
- ) from e
96
-
95
+ raise ValueError(msg) from e
97
96
  except Exception as e:
98
97
  msg = f"Unexpected error while loading the MuJoCo model: {e}"
99
- raise Exception(
100
- msg,
101
- )
98
+ raise Exception(msg)
102
99
 
103
100
  def _load_xml_file(self, xml: str, **kwargs) -> None:
104
101
  """Load a MuJoCo model from an XML file."""
@@ -119,6 +116,7 @@ class Wrapper:
119
116
 
120
117
  def _load_urdf_file(self, urdf_path: str, **kwargs: Any) -> None:
121
118
  """Process and load a URDF file for use with MuJoCo."""
119
+ # BUG: Generate actuators from joints in URDF so the controller can actually work
122
120
 
123
121
  def convert_dae_to_stl(meshdir: str) -> None:
124
122
  """Convert all DAE files in a directory (including subdirectories) to STL."""
@@ -241,13 +239,15 @@ class Wrapper:
241
239
  thread.join()
242
240
 
243
241
  @property
244
- def model(self) -> mujoco.MjModel:
242
+ def model(self) -> mjModel:
245
243
  """Read-only property to access the MjModel object."""
246
244
  return self._model
247
245
 
248
246
  @property
249
- def data(self) -> mujoco.MjData:
250
- """Read-only property to access the MjData single-step object."""
247
+ def data(self) -> mjData:
248
+ r"""Read-only property to access the MjData single-step object.\n
249
+ Use `captured_data` to access the entire simulation data.
250
+ """
251
251
  return self._data
252
252
 
253
253
  @property
@@ -272,9 +272,9 @@ class Wrapper:
272
272
 
273
273
  @frames.deleter
274
274
  def frames(self) -> None:
275
+ # BUG: This is not working as intended
275
276
  self._frames.clear()
276
277
  import gc
277
-
278
278
  gc.collect()
279
279
 
280
280
  @property
@@ -326,27 +326,32 @@ class Wrapper:
326
326
  self._model.vis.Global.offheight = self._height
327
327
 
328
328
  @property
329
- def init_conditions(self):
329
+ def initial_conditions(self):
330
330
  return self._initcond
331
331
 
332
- @init_conditions.setter
333
- def init_conditions(self, values) -> None:
332
+ @initial_conditions.setter
333
+ def initial_conditions(self, values) -> None:
334
334
  if not isinstance(values, dict):
335
335
  msg = "Initial conditions must be a dictionary."
336
336
  raise ValueError(msg)
337
- invalid_keys = [
338
- key for key in values if not hasattr(mujoco.MjData(self._model), key)
339
- ]
337
+
338
+ invalid_keys = [key for key in values if not hasattr(mujoco.MjData(self._model), key)]
340
339
  if invalid_keys:
341
340
  SimulationData._get_public_keys(self._data)
342
341
  msg = f"Invalid initial condition attributes: {', '.join(invalid_keys)}"
343
- raise ValueError(
344
- msg,
345
- )
342
+ raise ValueError(msg)
343
+
346
344
  self._initcond = values
347
345
 
346
+ # Apply initial conditions directly when set
347
+ for key, value in values.items():
348
+ if hasattr(self._data, key):
349
+ setattr(self._data, key, value)
350
+ else:
351
+ print_warning(f"'{key}' is not a valid attribute of MjData.")
352
+
348
353
  @property
349
- def controller(self) -> Optional[Callable[[mjModel, mjData, Any], None]]:
354
+ def controller(self) -> Callable[[mjModel, mjData, Any], None] | None:
350
355
  """Controller Function."""
351
356
  return self._controller
352
357
 
@@ -373,11 +378,11 @@ class Wrapper:
373
378
  try:
374
379
  return self._dr
375
380
  except AttributeError:
376
- print_warning(f"Use '{self.runSim.__name__}' first in order to access this value.")
381
+ print_warning(f"Use '{self.run.__name__}' first in order to access this value.")
377
382
  return None
378
383
 
379
384
  @data_rate.setter
380
- def data_rate(self, value) -> None:
385
+ def data_rate(self, value: int) -> None:
381
386
  if isinstance(value, float) and not isinstance(value, int):
382
387
  value = round(value)
383
388
  print_warning(f"Data rate must be an integer. Rounding to the nearest integer ({value}).")
@@ -396,20 +401,17 @@ class Wrapper:
396
401
  return self._model.opt.gravity
397
402
 
398
403
  @gravity.setter
399
- def gravity(self, values) -> None:
400
- if len(values) != 3:
404
+ def gravity(self, values: list | tuple | np.ndarray) -> None:
405
+ if not isinstance(values, (list, tuple, np.ndarray)) or len(values) != 3:
401
406
  msg = "Gravity must be a 3D vector."
402
407
  raise ValueError(msg)
403
- self._model.opt.gravity = values
408
+ self._model.opt.gravity = np.array(values)
404
409
 
405
- def _setInitialConditions(self) -> None:
406
- for key, value in self._initcond.items():
407
- if hasattr(self._data, key):
408
- setattr(self._data, key, value)
409
- else:
410
- print_warning(f"'{key}' is not a valid attribute of MjData.")
411
-
412
- def runSim(self, render: bool = False, camera: Optional[str] = None, data_rate: int = 100, interactive: bool = False, multi_thread: bool = False) -> "Wrapper":
410
+ def run(self,
411
+ render: bool = False,
412
+ camera: str | None = None,
413
+ interactive: bool = False,
414
+ multi_thread: bool = False) -> "Wrapper":
413
415
  """Run the simulation with optional rendering and controlled data capture.
414
416
 
415
417
  Args:
@@ -428,7 +430,6 @@ class Wrapper:
428
430
  try:
429
431
  mujoco.set_mjcb_control(self._controller) if self._controller else None
430
432
  mujoco.mj_resetData(self._model, self._data)
431
- self._setInitialConditions()
432
433
 
433
434
  # self._frames = [None] * (self.duration * self.fps + 1)
434
435
  frames = []
@@ -444,10 +445,7 @@ class Wrapper:
444
445
  dur = self._duration
445
446
  num_geoms = self._geom_names.__len__()
446
447
 
447
- # TODO: Fix references to data rate in future.
448
- # May be desired to specify in object creation
449
- self._dr = data_rate
450
- capture_rate = data_rate * self.ts
448
+ capture_rate = self.data_rate * self.ts
451
449
  capture_interval = max(1, int(1.0 / capture_rate))
452
450
  render_interval = max(1, int(1.0 / (self._fps * self.ts)))
453
451
 
@@ -461,17 +459,15 @@ class Wrapper:
461
459
 
462
460
  from . import COMPUTER, MAX_GEOM_SCALAR
463
461
  if COMPUTER.IDE == "jupyter":
464
- from tqdm.notebook import tqdm as bar
462
+ from tqdm.notebook import tqdm as ProgressBar
465
463
  else:
466
- from tqdm import tqdm as bar
467
-
468
- with (
469
- bar(total=total_steps,
470
- desc="Simulation",
471
- unit=" step", leave=False,
472
- ) as pbar,
473
- mujoco.Renderer(m, h, w, num_geoms * MAX_GEOM_SCALAR) as renderer,
474
- ):
464
+ from tqdm import tqdm as ProgressBar
465
+
466
+ # PEP Convention - Prevent Naming Conflict
467
+ _ProgressBar = ProgressBar(total=total_steps, desc="Simulation", unit=" step", leave=False)
468
+ _Renderer = mujoco.Renderer(m, h, w, num_geoms * MAX_GEOM_SCALAR)
469
+
470
+ with _ProgressBar as pbar, _Renderer as renderer:
475
471
  step = 0
476
472
  while d.time < dur:
477
473
  mj_step(m, d)
@@ -511,143 +507,196 @@ class Wrapper:
511
507
 
512
508
  return self
513
509
 
514
- def _window(self) -> None:
510
+ def _window(self, show_menu: bool =True) -> None:
515
511
  """Open a window to display the simulation in real time."""
512
+ try:
513
+ m = self._model
514
+ d = self._data
515
+
516
+ def key_callback(key) -> bool:
517
+ return key in (27, ord("q")) # 27 = ESC key, 'q' to quit
518
+
519
+ # NOTE: launch_passive may blocking despite docstring saying otherwise
520
+ _Viewer = mujoco.viewer.launch_passive(m, d,
521
+ show_left_ui=show_menu,
522
+ show_right_ui=show_menu,
523
+ key_callback=key_callback)
524
+ with _Viewer as viewer:
525
+ viewer.sync()
526
+ start_time = time.time()
516
527
 
528
+ try:
529
+ while viewer.is_running():
530
+ current_time = time.time()
531
+ dt = current_time - start_time # Time difference between frames
517
532
 
533
+ mujoco.mj_step(m, d) # Advance simulation by one step
534
+ viewer.sync() # Sync the viewer
535
+
536
+ start_time = current_time # Reset start_time for the next frame
537
+ time.sleep(max(0, 1.0 / self.fps - dt)) # Adjust sleep to match self.fps
538
+
539
+ except KeyboardInterrupt:
540
+ viewer.close()
541
+ except Exception as e:
542
+ msg = "An error occurred while running the simulation."
543
+ raise RuntimeError(msg) from e
544
+ finally:
545
+ mujoco.set_mjcb_control(None)
518
546
 
519
547
  def liveView(self, show_menu: bool = True) -> None:
520
548
  """Open a window to display the simulation in real time."""
521
- # BUG: Generate actuators from joints in URDF so the controller can actually work
549
+ # Run the window in a separate thread
550
+ gui = threading.Thread(target=self._window, kwargs={"show_menu": show_menu})
551
+ gui.start()
522
552
 
523
- def window() -> None:
524
- try:
525
- m = self._model
526
- d = self._data
527
-
528
- def key_callback(key) -> bool:
529
- if key in (27, ord("q")): # 27 = ESC key, 'q' to quit
530
- return True
531
- return False
532
-
533
- # NOTE: launch_passive is blocking despite docstring saying otherwise
534
- with mujoco.viewer.launch_passive(m, d,
535
- show_left_ui=show_menu,
536
- show_right_ui=show_menu,
537
- key_callback=key_callback) as viewer:
538
- viewer.sync()
539
- start_time = time.time()
540
-
541
- try:
542
- while viewer.is_running():
543
- current_time = time.time()
544
- dt = current_time - start_time # Time difference between frames
545
-
546
- mujoco.mj_step(m, d) # Advance simulation by one step
547
- viewer.sync() # Sync the viewer
548
-
549
- start_time = current_time # Reset start_time for the next frame
550
- # Sleep to match real-time simulation speed
551
- time.sleep(max(0, 0.01 - dt)) # Adjust sleep to match real-time
552
- except KeyboardInterrupt:
553
- viewer.close()
554
- except Exception as e:
555
- msg = "An error occurred while running the simulation."
556
- raise RuntimeError(msg) from e
557
- finally:
558
- mujoco.set_mjcb_control(None)
559
-
560
- viewer = threading.Thread(target=window)
561
- viewer.start()
562
-
563
- def renderFrame(self, t=0, frame=0, title=None) -> Optional[str]:
564
- """Render a specific frame as an image.
553
+ def _validate_and_extract_frames(
554
+ self,
555
+ frame_idx: int | tuple[int, int] | None = None,
556
+ time_idx: float | tuple[float, float] | None = None,
557
+ ) -> list[Any]:
558
+ """Validate and extract frames based on frame or time indices.
565
559
 
566
560
  Args:
567
- t (float): Time in seconds for which the frame should be rendered.
568
- frame (int): The frame index to render.
569
- title (str): The title of the rendered frame.
561
+ frame_idx (int or tuple, optional): Single frame index or (start, stop) frame indices.
562
+ time_idx (float or tuple, optional): Single time or (start, end) times in seconds.
570
563
 
571
564
  Returns:
572
- None
565
+ List of frames
566
+
567
+ Raises:
568
+ ValueError: For invalid input parameters.
573
569
 
574
570
  """
575
571
  if not self._frames:
576
572
  msg = "No frames captured to render."
577
573
  raise ValueError(msg)
578
574
 
579
- if t < 0 or frame < 0:
580
- msg = "Time and frame index must be greater than or equal to zero."
581
- raise ValueError(
582
- msg,
583
- )
584
-
585
- if frame and t:
586
- msg = "Can only specify singular time or frame parameter"
575
+ # Validate input parameters
576
+ if frame_idx is not None and time_idx is not None:
577
+ msg = "Can only specify either frame_idx or time_idx, not both."
587
578
  raise ValueError(msg)
588
579
 
589
- try:
590
- if t > 0:
591
- frame = self.t2f(t) # Convert time to frame index
580
+ # If both are None, use all frames
581
+ if frame_idx is None and time_idx is None:
582
+ return self._frames
583
+
584
+ # Handle time index conversion
585
+ if time_idx is not None:
586
+ if isinstance(time_idx, (int, float)):
587
+ frame_idx = self.t2f(time_idx)
588
+ elif isinstance(time_idx, tuple):
589
+ frame_idx = (self.t2f(time_idx[0]), self.t2f(time_idx[1]))
592
590
  else:
593
- frame = int(frame)
591
+ msg = "time_idx must be a number or a tuple of numbers."
592
+ raise ValueError(msg)
594
593
 
595
- plt.imshow(self._frames[frame])
596
- plt.axis("off")
597
- plt.title(title or f"Frame {frame}", loc="center")
598
- plt.show()
594
+ # Convert single index to tuple range
595
+ if isinstance(frame_idx, (int, float)):
596
+ frame_idx = (frame_idx, frame_idx + 1)
599
597
 
600
- except IndexError as e:
601
- msg = f"Invalid frame index: {e}"
598
+ # Validate frame indices
599
+ start, stop = frame_idx
600
+ max_frames = len(self._frames)
601
+
602
+ if start < 0:
603
+ msg = f"Start index must be non-negative. Got {start}."
602
604
  raise ValueError(msg)
603
- except TypeError as e:
604
- msg = f"Invalid type for time or frame: {e}"
605
+ if stop > max_frames:
606
+ msg = f"Stop index must not exceed total frames ({max_frames}). Got {stop}."
605
607
  raise ValueError(msg)
606
- except Exception as e:
607
- msg = f"Unexpected error while rendering frame: {e}"
608
- raise RuntimeError(msg)
608
+ if start >= stop:
609
+ msg = f"Start index ({start}) must be less than stop index ({stop})."
610
+ raise ValueError(msg)
611
+
612
+ # Select subset of frames
613
+ return self._frames[start:stop]
609
614
 
610
- def renderMedia(self, codec="gif", title=None, save=False) -> media:
611
- """Render the simulation as a video or GIF, with an option to save to a file.
615
+ def show(
616
+ self,
617
+ title: str | None = None,
618
+ codec: str = "gif",
619
+ frame_idx: int | tuple[int, int] | None = None,
620
+ time_idx: float | tuple[float, float] | None = None,
621
+ ) -> None:
622
+ """Render specific frame(s) as a video or GIF in a window.
612
623
 
613
624
  Args:
614
- codec (str): The media format to use ("gif" or "mp4").
615
- title (str): The filename or window title for the media.
616
- save (bool): Whether to save the media to a file.
625
+ title (str, optional): Title for the rendered media.
626
+ codec (str, optional): Video codec/format. Defaults to "gif".
627
+ frame_idx (int or tuple, optional): Single frame index or (start, stop) frame indices.
628
+ time_idx (float or tuple, optional): Single time or (start, end) times in seconds.
629
+
630
+ Raises:
631
+ ValueError: If no frames are captured or invalid input parameters.
617
632
 
618
633
  """
619
- if not self._frames:
620
- msg = "No frames captured to create media."
621
- raise ValueError(msg)
622
- # Display the media in a window
623
- if not save:
634
+ try:
635
+ # Extract frames
636
+ subset_frames = self._validate_and_extract_frames(
637
+ frame_idx=frame_idx,
638
+ time_idx=time_idx,
639
+ )
640
+
641
+ # Show the video
624
642
  media.show_video(
625
- self._frames,
626
- fps=self._fps,
643
+ subset_frames,
644
+ fps=1 if len(subset_frames) == 1 else self._fps,
627
645
  width=self._width,
628
646
  height=self._height,
629
647
  codec=codec,
630
648
  title=title,
631
649
  )
632
- return None
650
+ except Exception as e:
651
+ msg = "Error while showing video subset."
652
+ raise Exception(msg) from e
653
+
654
+ def save(
655
+ self,
656
+ title: str = "render",
657
+ codec: str = "gif",
658
+ frame_idx: int | tuple[int, int] | None = None,
659
+ time_idx: float | tuple[float, float] | None = None,
660
+ ) -> str:
661
+ """Save specific frame(s) as a video or GIF to a file.
662
+
663
+ Args:
664
+ title (str, optional): Filename for the saved media.
665
+ codec (str, optional): Video codec/format. Defaults to "gif".
666
+ frame_idx (int or tuple, optional): Single frame index or (start, stop) frame indices.
667
+ time_idx (float or tuple, optional): Single time or (start, end) times in seconds.
668
+
669
+ Returns:
670
+ str: Absolute path to the saved file.
671
+
672
+ Raises:
673
+ ValueError: If no frames are captured or invalid input parameters.
633
674
 
634
- if not title.endswith(f".{codec}"):
635
- title += f".{codec}"
675
+ """
676
+ try:
677
+ # Extract frames
678
+ subset_frames = self._validate_and_extract_frames(
679
+ frame_idx=frame_idx,
680
+ time_idx=time_idx,
681
+ )
636
682
 
637
- # Save the frames to the specified file
638
- available_codecs = ["gif", "mp4", "h264", "hevc", "vp9"]
639
- if codec in available_codecs:
683
+ # Ensure the title ends with the correct codec extension
684
+ if not title.endswith(f".{codec}"):
685
+ title += f".{codec}"
686
+
687
+ # Save the video
640
688
  media.write_video(
641
- title if not None else "render",
642
- self._frames,
643
- fps=self._fps,
689
+ title,
690
+ subset_frames,
691
+ fps=1 if len(subset_frames) == 1 else self._fps,
644
692
  codec=codec,
645
693
  )
646
- else:
647
- msg = f"Unsupported codec '{codec}'. Supported codecs are {', '.join(available_codecs)}"
648
- raise ValueError(msg)
649
694
 
650
- return os.path.abspath(title)
695
+ return os.path.abspath(title)
696
+
697
+ except Exception as e:
698
+ msg = "Error while saving video subset."
699
+ raise Exception(msg) from e
651
700
 
652
701
  @lru_cache(maxsize=100)
653
702
  def t2f(self, t: float) -> int:
@@ -661,7 +710,7 @@ class Wrapper:
661
710
  """Convert frame index to time."""
662
711
  return frame / self._fps
663
712
 
664
- def getBodyData(self, body_name: str, data_name: Optional[str] = None) -> np.ndarray:
713
+ def getBodyData(self, body_name: str, data_name: str | None = None) -> np.ndarray:
665
714
  """Get the data for a specific body in the simulation.
666
715
 
667
716
  Args:
@@ -2,14 +2,14 @@
2
2
  requires = ["poetry-core"]
3
3
  build-backend = "poetry.core.masonry.api"
4
4
 
5
- [tool.poetry]
5
+ [project]
6
6
  name = "mujoco-toolbox"
7
- version = "0.2.2"
7
+ version = "0.3.0"
8
8
  description = "A modern package to handle MuJoCo environments"
9
9
  readme = "README.md"
10
10
  license = "MIT"
11
11
  packages = [{ include = "mujoco_toolbox" }]
12
- authors = ["MGross21 <MGross21@TODO.com>"]
12
+ authors = [{ name = "MGross21", email = "MGross21@TODO.com" }]
13
13
  keywords = ["MuJoCo", "robotics", "simulation", "toolbox", "wrapper"]
14
14
  classifiers = [
15
15
  "Development Status :: 3 - Alpha",
@@ -18,31 +18,35 @@ classifiers = [
18
18
  "Programming Language :: Python :: 3.10",
19
19
  "Topic :: Software Development :: Libraries"
20
20
  ]
21
+ requires-python = ">=3.10"
22
+ dependencies = [
23
+ "mujoco",
24
+ "numpy",
25
+ "scipy",
26
+ "matplotlib",
27
+ "pyyaml",
28
+ "mediapy",
29
+ "tqdm",
30
+ "ipywidgets",
31
+ "ffmpeg",
32
+ "trimesh",
33
+ "screeninfo",
34
+ "colorama"
35
+ ]
36
+
37
+ # [project.extras]
38
+ # beautified = ["colorama","ipywidgets","tqdm"]
39
+ # gym = ["gym"]
21
40
 
22
- [tool.poetry.urls]
41
+ [project.urls]
23
42
  homepage = "https://github.com/MGross21/mujoco-toolbox"
24
43
  repository = "https://github.com/MGross21/mujoco-toolbox"
25
44
  documentation = "https://github.com/MGross21/mujoco-toolbox"
26
45
  "Bug Tracker" = "https://github.com/MGross21/mujoco-toolbox/issues"
27
46
 
28
- [tool.poetry.dependencies]
29
- python = ">=3.10"
30
- mujoco = "*"
31
- numpy = "*"
32
- scipy = "*"
33
- matplotlib = "*"
34
- pyyaml = "*"
35
- mediapy = "*"
36
- tqdm = "*"
37
- ipywidgets = "*"
38
- ffmpeg = "*"
39
- trimesh = "*"
40
- screeninfo = "*"
41
- colorama = "*"
42
-
43
47
  [tool.poetry.group.dev.dependencies]
44
48
  pytest = "*"
45
- ruff = "^0.9.9"
49
+ ruff = "^0.11.2"
46
50
  mypy = "^1.15.0"
47
51
  pytypes = "^1.0b10"
48
52
  black = "^25.1.0"
@@ -57,11 +61,11 @@ version_variables = [
57
61
  "pyproject.toml:tool.poetry.version",
58
62
  "CITATION.cff:version"
59
63
  ]
60
- branch = "main"
64
+ branch = "main" # Branch to track for triggering the release process
61
65
  changelog_file = "CHANGELOG.md"
62
66
  build_command = "pip install poetry && poetry build"
63
- dist_path = "dist/"
64
- create_pr = true
67
+ dist_path = "dist/" # Directory where the distribution files will be stored
68
+ create_pr = true # Option to create a pull request when a new release version is created
65
69
  pr_branch_prefix = "release/"
66
70
  upload_to_pypi = true
67
71
  upload_to_release = true
@@ -72,8 +76,12 @@ commit_message = "chore(release): {version} [skip ci]"
72
76
  commit_author = "semantic-release <semantic-release@example.com>"
73
77
 
74
78
  [tool.semantic_release.plugins]
75
- commit_analyzer = "@semantic-release/commit-analyzer"
76
- release_notes_generator = "@semantic-release/release-notes-generator"
79
+ commit_analyzer = "@semantic-release/commit-analyzer" # Analyzes commit messages to determine version bumps
80
+ release_notes_generator = "@semantic-release/release-notes-generator" # Generates release notes based on commit messages
81
+ changelog = "@semantic-release/changelog" # Handles changelog generation
82
+ git = "@semantic-release/git" # Commits changes (version bump, changelog) back to the repo
83
+ github = "@semantic-release/github" # Creates GitHub releases (tags) and adds release notes
84
+ pypi = "@semantic-release/pypi" # Uploads the distribution to PyPI
77
85
 
78
86
  [tool.semantic_release.commit_analyzer]
79
87
  preset = "angular"
@@ -111,7 +119,7 @@ preset_config = {types = [
111
119
  ]}
112
120
 
113
121
  [tool.ruff]
114
- line-length = 120
122
+ line-length = 79 # PEP8 line length
115
123
 
116
124
  [tool.ruff.lint]
117
125
  # select = ["E", "F", "W", "C90"] # Select error, warning, and flake8 rules
File without changes