mujoco-toolbox 0.3.4__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.
- mujoco_toolbox-0.3.4/LICENSE +21 -0
- mujoco_toolbox-0.3.4/PKG-INFO +150 -0
- mujoco_toolbox-0.3.4/README.md +96 -0
- mujoco_toolbox-0.3.4/mujoco_toolbox/__init__.py +77 -0
- mujoco_toolbox-0.3.4/mujoco_toolbox/assets.py +29 -0
- mujoco_toolbox-0.3.4/mujoco_toolbox/builder.py +229 -0
- mujoco_toolbox-0.3.4/mujoco_toolbox/controller.py +164 -0
- mujoco_toolbox-0.3.4/mujoco_toolbox/utils.py +134 -0
- mujoco_toolbox-0.3.4/mujoco_toolbox/wrapper.py +893 -0
- mujoco_toolbox-0.3.4/pyproject.toml +150 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Michael Gross
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,150 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: mujoco-toolbox
|
|
3
|
+
Version: 0.3.4
|
|
4
|
+
Summary: A modern package to handle MuJoCo environments
|
|
5
|
+
License: MIT License
|
|
6
|
+
|
|
7
|
+
Copyright (c) 2025 Michael Gross
|
|
8
|
+
|
|
9
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
10
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
11
|
+
in the Software without restriction, including without limitation the rights
|
|
12
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
13
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
14
|
+
furnished to do so, subject to the following conditions:
|
|
15
|
+
|
|
16
|
+
The above copyright notice and this permission notice shall be included in all
|
|
17
|
+
copies or substantial portions of the Software.
|
|
18
|
+
|
|
19
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
20
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
21
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
22
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
23
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
24
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
25
|
+
SOFTWARE.
|
|
26
|
+
Keywords: MuJoCo,mujoco,robotics,simulation,toolbox,wrapper
|
|
27
|
+
Author: MGross21
|
|
28
|
+
Author-email: MGross21@TODO.com
|
|
29
|
+
Requires-Python: >=3.10
|
|
30
|
+
Classifier: Development Status :: 4 - Beta
|
|
31
|
+
Classifier: Intended Audience :: Developers
|
|
32
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
33
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
34
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
35
|
+
Requires-Dist: colorama
|
|
36
|
+
Requires-Dist: ffmpeg
|
|
37
|
+
Requires-Dist: ipywidgets
|
|
38
|
+
Requires-Dist: matplotlib
|
|
39
|
+
Requires-Dist: mediapy
|
|
40
|
+
Requires-Dist: mujoco
|
|
41
|
+
Requires-Dist: numpy
|
|
42
|
+
Requires-Dist: pyyaml
|
|
43
|
+
Requires-Dist: scipy
|
|
44
|
+
Requires-Dist: screeninfo
|
|
45
|
+
Requires-Dist: tqdm
|
|
46
|
+
Requires-Dist: trimesh
|
|
47
|
+
Project-URL: Bug Tracker, https://github.com/MGross21/mujoco-toolbox/issues
|
|
48
|
+
Project-URL: Documentation, https://github.com/MGross21/mujoco-toolbox
|
|
49
|
+
Project-URL: Homepage, https://github.com/MGross21/mujoco-toolbox
|
|
50
|
+
Project-URL: Repository, https://github.com/MGross21/mujoco-toolbox
|
|
51
|
+
Project-URL: changelog, https://github.com/MGross21/mujoco-toolbox/blob/main/CHANGELOG.md
|
|
52
|
+
Description-Content-Type: text/markdown
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
[](https://github.com/MGross21/mujoco-toolbox/actions/workflows/publish.yml)
|
|
58
|
+
[](https://github.com/MGross21/mujoco-toolbox/actions/workflows/docs.yml)
|
|
59
|
+
|
|
60
|
+
# Mujoco Toolbox
|
|
61
|
+
|
|
62
|
+
Streamlines the MuJoCo Physics Simulator
|
|
63
|
+
|
|
64
|
+
## PyPI Package
|
|
65
|
+
|
|
66
|
+
```bash
|
|
67
|
+
pip install -U mujoco-toolbox
|
|
68
|
+
```
|
|
69
|
+
|
|
70
|
+
## Local Install
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
pip install -U git+https://github.com/MGross21/mujoco-toolbox.git@main
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
## Extra Packages
|
|
77
|
+
|
|
78
|
+
<details>
|
|
79
|
+
<summary><b>FFMPEG</b></summary>
|
|
80
|
+
|
|
81
|
+
</br>
|
|
82
|
+
|
|
83
|
+
*Required for [mediapy](https://google.github.io/mediapy/mediapy.html) dependency*
|
|
84
|
+
|
|
85
|
+
**Windows**
|
|
86
|
+
|
|
87
|
+
```bash
|
|
88
|
+
winget install ffmpeg
|
|
89
|
+
ffmpeg -version
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
**Linux**
|
|
93
|
+
|
|
94
|
+
```bash
|
|
95
|
+
sudo apt update && sudo apt install ffmpeg
|
|
96
|
+
ffmpeg -version
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
**MacOS**
|
|
100
|
+
|
|
101
|
+
*Using Homebrew*
|
|
102
|
+
|
|
103
|
+
```bash
|
|
104
|
+
brew install ffmpeg
|
|
105
|
+
ffmpeg -version
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
*Using MacPorts*
|
|
109
|
+
|
|
110
|
+
```bash
|
|
111
|
+
sudo port install ffmpeg
|
|
112
|
+
ffmpeg -version
|
|
113
|
+
```
|
|
114
|
+
|
|
115
|
+
</details>
|
|
116
|
+
|
|
117
|
+
## Example Script
|
|
118
|
+
|
|
119
|
+
*Bare minimum to run MuJoCo simulation and display result*
|
|
120
|
+
|
|
121
|
+
```python
|
|
122
|
+
import mujoco_toolbox as mjtb
|
|
123
|
+
|
|
124
|
+
mjtb.Wrapper("path/to/your/xml").run(render=True).save()
|
|
125
|
+
```
|
|
126
|
+
|
|
127
|
+
## Pre-Made Controllers
|
|
128
|
+
|
|
129
|
+
```python
|
|
130
|
+
from mujoco_toolbox.controller import (
|
|
131
|
+
cosine_controller,
|
|
132
|
+
random_controller,
|
|
133
|
+
live_controller,
|
|
134
|
+
sine_controller,
|
|
135
|
+
step_controller,
|
|
136
|
+
)
|
|
137
|
+
|
|
138
|
+
# Wrapper can use custom controllers as well!
|
|
139
|
+
```
|
|
140
|
+
|
|
141
|
+
## File Support
|
|
142
|
+
|
|
143
|
+
### XML / MJCF (Native)
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
|
|
147
|
+
### URDF
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
|
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
[](https://github.com/MGross21/mujoco-toolbox/actions/workflows/publish.yml)
|
|
5
|
+
[](https://github.com/MGross21/mujoco-toolbox/actions/workflows/docs.yml)
|
|
6
|
+
|
|
7
|
+
# Mujoco Toolbox
|
|
8
|
+
|
|
9
|
+
Streamlines the MuJoCo Physics Simulator
|
|
10
|
+
|
|
11
|
+
## PyPI Package
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
pip install -U mujoco-toolbox
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
## Local Install
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
pip install -U git+https://github.com/MGross21/mujoco-toolbox.git@main
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
## Extra Packages
|
|
24
|
+
|
|
25
|
+
<details>
|
|
26
|
+
<summary><b>FFMPEG</b></summary>
|
|
27
|
+
|
|
28
|
+
</br>
|
|
29
|
+
|
|
30
|
+
*Required for [mediapy](https://google.github.io/mediapy/mediapy.html) dependency*
|
|
31
|
+
|
|
32
|
+
**Windows**
|
|
33
|
+
|
|
34
|
+
```bash
|
|
35
|
+
winget install ffmpeg
|
|
36
|
+
ffmpeg -version
|
|
37
|
+
```
|
|
38
|
+
|
|
39
|
+
**Linux**
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
sudo apt update && sudo apt install ffmpeg
|
|
43
|
+
ffmpeg -version
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
**MacOS**
|
|
47
|
+
|
|
48
|
+
*Using Homebrew*
|
|
49
|
+
|
|
50
|
+
```bash
|
|
51
|
+
brew install ffmpeg
|
|
52
|
+
ffmpeg -version
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
*Using MacPorts*
|
|
56
|
+
|
|
57
|
+
```bash
|
|
58
|
+
sudo port install ffmpeg
|
|
59
|
+
ffmpeg -version
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
</details>
|
|
63
|
+
|
|
64
|
+
## Example Script
|
|
65
|
+
|
|
66
|
+
*Bare minimum to run MuJoCo simulation and display result*
|
|
67
|
+
|
|
68
|
+
```python
|
|
69
|
+
import mujoco_toolbox as mjtb
|
|
70
|
+
|
|
71
|
+
mjtb.Wrapper("path/to/your/xml").run(render=True).save()
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Pre-Made Controllers
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
from mujoco_toolbox.controller import (
|
|
78
|
+
cosine_controller,
|
|
79
|
+
random_controller,
|
|
80
|
+
live_controller,
|
|
81
|
+
sine_controller,
|
|
82
|
+
step_controller,
|
|
83
|
+
)
|
|
84
|
+
|
|
85
|
+
# Wrapper can use custom controllers as well!
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
## File Support
|
|
89
|
+
|
|
90
|
+
### XML / MJCF (Native)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
### URDF
|
|
95
|
+
|
|
96
|
+
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
"""Mujoco Toolbox.
|
|
2
|
+
==============
|
|
3
|
+
A toolbox for working with MuJoCo simulations. This package provides various utilities and controllers to facilitate
|
|
4
|
+
the simulation process.
|
|
5
|
+
|
|
6
|
+
Modules
|
|
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.
|
|
11
|
+
- `assets`: Contains pre-defined assets for building MuJoCo models.
|
|
12
|
+
|
|
13
|
+
Constants
|
|
14
|
+
----------
|
|
15
|
+
- CAPTURE_PARAMETERS: List of MjData fields to capture during simulation.
|
|
16
|
+
|
|
17
|
+
Notes
|
|
18
|
+
-----
|
|
19
|
+
This package is still under development. Report any issues to https://github.com/MGross21/mujoco-toolbox/issues.
|
|
20
|
+
|
|
21
|
+
"""
|
|
22
|
+
|
|
23
|
+
from .assets import WORLD_ASSETS, glovebox
|
|
24
|
+
from .builder import Builder
|
|
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
|
|
33
|
+
from .wrapper import Wrapper
|
|
34
|
+
|
|
35
|
+
__version__ = "0.3.4"
|
|
36
|
+
__author__ = "Michael Gross"
|
|
37
|
+
__github_repo__ = "mujoco-toolbox"
|
|
38
|
+
__license__ = "MIT"
|
|
39
|
+
__status__ = "Development"
|
|
40
|
+
__all__ = [
|
|
41
|
+
"CAPTURE_PARAMETERS",
|
|
42
|
+
"COMPUTER",
|
|
43
|
+
"WORLD_ASSETS",
|
|
44
|
+
"Builder",
|
|
45
|
+
"Wrapper",
|
|
46
|
+
"cosine_controller",
|
|
47
|
+
"glovebox",
|
|
48
|
+
"live_controller",
|
|
49
|
+
"random_controller",
|
|
50
|
+
"sine_controller",
|
|
51
|
+
"step_controller",
|
|
52
|
+
]
|
|
53
|
+
|
|
54
|
+
COMPUTER = _Platform() # Singleton Instance
|
|
55
|
+
MAX_GEOM_SCALAR: int = 2 # Scalar value for mujoco.Renderer.max_geom
|
|
56
|
+
PROGRESS_BAR: bool = False # Enable/Disable progress bar
|
|
57
|
+
CAPTURE_PARAMETERS = [ # MjData default fields to capture during simulation
|
|
58
|
+
"time",
|
|
59
|
+
"qpos",
|
|
60
|
+
"qvel",
|
|
61
|
+
"act",
|
|
62
|
+
"qacc",
|
|
63
|
+
"xpos",
|
|
64
|
+
"xquat",
|
|
65
|
+
"xmat",
|
|
66
|
+
"ctrl",
|
|
67
|
+
"sensordata",
|
|
68
|
+
]
|
|
69
|
+
|
|
70
|
+
if __version__.startswith("0"):
|
|
71
|
+
from .utils import _print_warning
|
|
72
|
+
_print_warning(
|
|
73
|
+
f"{__package__} (v{__version__}) is still under development.",
|
|
74
|
+
f"Report any issues to https://github.com/MGross21/{__github_repo__}/issues",
|
|
75
|
+
)
|
|
76
|
+
del _print_warning
|
|
77
|
+
del _Platform
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
WORLD_ASSETS = """
|
|
2
|
+
<asset>
|
|
3
|
+
<texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="32" height="512"/>
|
|
4
|
+
<texture name="body" type="cube" builtin="flat" mark="cross" width="128" height="128" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1"/>
|
|
5
|
+
<material name="body" texture="body" texuniform="true" rgba="0.8 0.6 .4 1"/>
|
|
6
|
+
<texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
|
|
7
|
+
<material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
|
|
8
|
+
<texture name="background" type="2d" builtin="flat" width="256" height="256" rgb1="1 1 1"/>
|
|
9
|
+
<material name="background" texture="background" texuniform="true" rgba="1 1 1 1"/>
|
|
10
|
+
</asset>
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
def glovebox(width=1.25, depth=0.75, height=1.0, glass_thickness=0.05) -> str:
|
|
14
|
+
return f"""
|
|
15
|
+
<mujoco>
|
|
16
|
+
<asset>
|
|
17
|
+
<material name="glass" rgba="1 1 1 0.2"/>
|
|
18
|
+
</asset>
|
|
19
|
+
<worldbody>
|
|
20
|
+
<body name="walls" pos="0 0 0">
|
|
21
|
+
<geom type="box" size="{width/2} {glass_thickness/2} {height/2}" material="glass" pos="0 {depth/2+glass_thickness/2} {height/2}"/>
|
|
22
|
+
<geom type="box" size="{width/2} {glass_thickness/2} {height/2}" material="glass" pos="0 {-depth/2-glass_thickness/2} {height/2}"/>
|
|
23
|
+
<geom type="box" size="{glass_thickness/2} {depth/2} {height/2}" material="glass" pos="{width/2-glass_thickness/2} 0 {height/2}"/>
|
|
24
|
+
<geom type="box" size="{glass_thickness/2} {depth/2} {height/2}" material="glass" pos="{-width/2+glass_thickness/2} 0 {height/2}"/>
|
|
25
|
+
<geom type="plane" size="{width/2} {depth/2+glass_thickness} {glass_thickness/2}" material="glass" pos="0 0 {height}"/>
|
|
26
|
+
</body>
|
|
27
|
+
</worldbody>
|
|
28
|
+
</mujoco>
|
|
29
|
+
"""
|
|
@@ -0,0 +1,229 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import xml.etree.ElementTree as ET
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class Builder:
|
|
6
|
+
"""Builder class for loading, merging, and saving MuJoCo models from XML strings or file paths.
|
|
7
|
+
|
|
8
|
+
Attributes:
|
|
9
|
+
xml (str): XML string of the model.
|
|
10
|
+
|
|
11
|
+
Examples:
|
|
12
|
+
>>> # Load a model
|
|
13
|
+
>>> builder1 = Builder("path/to/model1.xml")
|
|
14
|
+
>>> # Merge two models from two object instances
|
|
15
|
+
>>> builder2 = Builder("path/to/model2.xml")
|
|
16
|
+
>>> merged_builder = builder1 + builder2
|
|
17
|
+
>>> merged_builder.save("path/to/merged_model.xml")
|
|
18
|
+
>>> # Merge two models from two file paths
|
|
19
|
+
>>> builder = Builder("path/to/model1.xml", "path/to/model2.xml")
|
|
20
|
+
>>> builder.save("path/to/merged_model.xml")
|
|
21
|
+
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
def __init__(self, *args: str) -> None:
|
|
25
|
+
"""Initialize with an XML string or a file path. Input is required."""
|
|
26
|
+
for arg in args:
|
|
27
|
+
if not isinstance(arg, str):
|
|
28
|
+
msg = "Input must be an XML string or a file path"
|
|
29
|
+
raise TypeError(msg)
|
|
30
|
+
if len(args) == 0:
|
|
31
|
+
msg = "Input is required to initialize the Builder"
|
|
32
|
+
raise ValueError(msg)
|
|
33
|
+
if len(args) > 1:
|
|
34
|
+
# Create a new Builder instance for each additional argument and merge them
|
|
35
|
+
self.tree, self.root = self._load_model(args[0])
|
|
36
|
+
for additional_arg in args[1:]:
|
|
37
|
+
additional_builder = Builder(additional_arg)
|
|
38
|
+
self.__add__(additional_builder)
|
|
39
|
+
else:
|
|
40
|
+
xml_input = args[0]
|
|
41
|
+
self.tree, self.root = self._load_model(xml_input)
|
|
42
|
+
|
|
43
|
+
def _load_model(self, xml_input):
|
|
44
|
+
"""Load a MuJoCo model from a file or XML string."""
|
|
45
|
+
if not isinstance(xml_input, str):
|
|
46
|
+
msg = "Input must be either an XML string or a file path"
|
|
47
|
+
raise ValueError(msg)
|
|
48
|
+
|
|
49
|
+
# Check if input is an XML string
|
|
50
|
+
if xml_input.strip().startswith("<"):
|
|
51
|
+
# Try parsing as XML string
|
|
52
|
+
try:
|
|
53
|
+
root = ET.fromstring(xml_input)
|
|
54
|
+
tree = ET.ElementTree(root)
|
|
55
|
+
|
|
56
|
+
# If it's a URDF (robot tag), wrap it in mujoco
|
|
57
|
+
if root.tag == "robot":
|
|
58
|
+
return self._wrap_urdf_in_mujoco(tree)
|
|
59
|
+
|
|
60
|
+
# If root is not mujoco, wrap it
|
|
61
|
+
if root.tag != "mujoco":
|
|
62
|
+
mujoco_elem = ET.Element("mujoco")
|
|
63
|
+
mujoco_elem.append(root)
|
|
64
|
+
tree = ET.ElementTree(mujoco_elem)
|
|
65
|
+
|
|
66
|
+
return tree, tree.getroot()
|
|
67
|
+
except ET.ParseError:
|
|
68
|
+
msg = "Invalid XML string provided"
|
|
69
|
+
raise ValueError(msg)
|
|
70
|
+
else:
|
|
71
|
+
# Treat input as file path
|
|
72
|
+
if not os.path.exists(xml_input):
|
|
73
|
+
msg = f"File not found: {xml_input}"
|
|
74
|
+
raise FileNotFoundError(msg)
|
|
75
|
+
|
|
76
|
+
# Handle URDF files
|
|
77
|
+
if xml_input.lower().endswith(".urdf"):
|
|
78
|
+
return self._wrap_urdf_in_mujoco(ET.parse(xml_input))
|
|
79
|
+
|
|
80
|
+
# Parse regular XML file
|
|
81
|
+
try:
|
|
82
|
+
tree = ET.parse(xml_input)
|
|
83
|
+
root = tree.getroot()
|
|
84
|
+
|
|
85
|
+
# If root is not mujoco, wrap it
|
|
86
|
+
if root.tag != "mujoco":
|
|
87
|
+
mujoco_elem = ET.Element("mujoco")
|
|
88
|
+
mujoco_elem.append(root)
|
|
89
|
+
tree = ET.ElementTree(mujoco_elem)
|
|
90
|
+
|
|
91
|
+
return tree, tree.getroot()
|
|
92
|
+
except ET.ParseError:
|
|
93
|
+
msg = f"Invalid XML file: {xml_input}"
|
|
94
|
+
raise ValueError(msg)
|
|
95
|
+
|
|
96
|
+
def _wrap_urdf_in_mujoco(self, urdf_input):
|
|
97
|
+
"""Wrap a URDF inside <mujoco> tags for MuJoCo compatibility.
|
|
98
|
+
Input can be a file path, ElementTree, or XML string.
|
|
99
|
+
"""
|
|
100
|
+
# Handle string input (could be a file path or XML string)
|
|
101
|
+
if isinstance(urdf_input, str):
|
|
102
|
+
# Check if it's an XML string
|
|
103
|
+
if urdf_input.strip().startswith("<"):
|
|
104
|
+
try:
|
|
105
|
+
root = ET.fromstring(urdf_input)
|
|
106
|
+
tree = ET.ElementTree(root)
|
|
107
|
+
except ET.ParseError:
|
|
108
|
+
msg = "Invalid URDF XML string provided"
|
|
109
|
+
raise ValueError(msg)
|
|
110
|
+
else:
|
|
111
|
+
# It's a file path
|
|
112
|
+
try:
|
|
113
|
+
tree = ET.parse(urdf_input)
|
|
114
|
+
except:
|
|
115
|
+
msg = f"Could not parse URDF file: {urdf_input}"
|
|
116
|
+
raise ValueError(msg)
|
|
117
|
+
# Handle ElementTree input
|
|
118
|
+
elif isinstance(urdf_input, ET.ElementTree):
|
|
119
|
+
tree = urdf_input
|
|
120
|
+
else:
|
|
121
|
+
msg = "URDF input must be a string or ElementTree"
|
|
122
|
+
raise TypeError(msg)
|
|
123
|
+
|
|
124
|
+
root = tree.getroot()
|
|
125
|
+
|
|
126
|
+
# Handle the case where the root might already be a mujoco element
|
|
127
|
+
if root.tag == "mujoco":
|
|
128
|
+
return tree, root
|
|
129
|
+
|
|
130
|
+
# Create mujoco element and append the URDF root
|
|
131
|
+
mujoco_elem = ET.Element("mujoco")
|
|
132
|
+
mujoco_elem.append(root)
|
|
133
|
+
|
|
134
|
+
# Create a new tree with the <mujoco> root
|
|
135
|
+
return ET.ElementTree(mujoco_elem), mujoco_elem
|
|
136
|
+
|
|
137
|
+
def _merge_tags(self, tag_name, root_1, root_2) -> None:
|
|
138
|
+
"""Merge a specific tag (e.g., worldbody, asset) from two models."""
|
|
139
|
+
section_1 = root_1.find(tag_name)
|
|
140
|
+
section_2 = root_2.find(tag_name)
|
|
141
|
+
|
|
142
|
+
# If target section doesn't exist in model 1 but exists in model 2, create it
|
|
143
|
+
if section_1 is None and section_2 is not None:
|
|
144
|
+
section_1 = ET.SubElement(root_1, tag_name)
|
|
145
|
+
|
|
146
|
+
# Only proceed if both sections exist now
|
|
147
|
+
if section_1 is not None and section_2 is not None:
|
|
148
|
+
# Merge the contents by appending all elements from model 2 to model 1
|
|
149
|
+
for element in list(section_2): # Use list() to avoid modification during iteration
|
|
150
|
+
section_1.append(element)
|
|
151
|
+
|
|
152
|
+
def __add__(self, other):
|
|
153
|
+
"""Implement the + operator for merging two models."""
|
|
154
|
+
model_root = None
|
|
155
|
+
|
|
156
|
+
if isinstance(other, Builder):
|
|
157
|
+
# Merge the current model with another Builder instance
|
|
158
|
+
model_root = other.root
|
|
159
|
+
elif isinstance(other, str):
|
|
160
|
+
# If the input is a string, load as XML
|
|
161
|
+
_, model_root = self._load_model(other)
|
|
162
|
+
else:
|
|
163
|
+
msg = "Addition only supported between Builder objects or with XML strings/files"
|
|
164
|
+
raise TypeError(msg)
|
|
165
|
+
|
|
166
|
+
# List of sections to merge
|
|
167
|
+
sections_to_merge = ["asset", "worldbody", "camera", "light", "contact", "equality",
|
|
168
|
+
"sensor", "actuator", "default", "tendon", "include"]
|
|
169
|
+
|
|
170
|
+
# Merge all relevant sections
|
|
171
|
+
for section in sections_to_merge:
|
|
172
|
+
self._merge_tags(section, self.root, model_root)
|
|
173
|
+
|
|
174
|
+
# Return self to allow for method chaining
|
|
175
|
+
return self
|
|
176
|
+
|
|
177
|
+
def __radd__(self, other):
|
|
178
|
+
"""Implement the reverse + operator for merging two models."""
|
|
179
|
+
if other == 0: # Handle sum() starting value
|
|
180
|
+
return self
|
|
181
|
+
return self.__add__(other)
|
|
182
|
+
|
|
183
|
+
def save(self, file_path) -> "os.path":
|
|
184
|
+
"""Save the merged model to a file."""
|
|
185
|
+
if self.tree is not None:
|
|
186
|
+
# Format the XML with proper indentation before saving
|
|
187
|
+
self._indent_xml(self.root)
|
|
188
|
+
self.tree.write(file_path, encoding="utf-8", xml_declaration=True)
|
|
189
|
+
else:
|
|
190
|
+
msg = "No model loaded. Cannot save."
|
|
191
|
+
raise ValueError(msg)
|
|
192
|
+
return os.path.abspath(file_path)
|
|
193
|
+
|
|
194
|
+
def _indent_xml(self, elem, level=0) -> None:
|
|
195
|
+
"""Add proper indentation to make the XML file more readable."""
|
|
196
|
+
i = "\n" + level * " "
|
|
197
|
+
if len(elem):
|
|
198
|
+
if not elem.text or not elem.text.strip():
|
|
199
|
+
elem.text = i + " "
|
|
200
|
+
if not elem.tail or not elem.tail.strip():
|
|
201
|
+
elem.tail = i
|
|
202
|
+
for sub_elem in elem:
|
|
203
|
+
self._indent_xml(sub_elem, level + 1)
|
|
204
|
+
if not elem[-1].tail or not elem[-1].tail.strip():
|
|
205
|
+
elem[-1].tail = i
|
|
206
|
+
elif level and (not elem.tail or not elem.tail.strip()):
|
|
207
|
+
elem.tail = i
|
|
208
|
+
|
|
209
|
+
def __str__(self) -> str:
|
|
210
|
+
"""Return the XML string of the model."""
|
|
211
|
+
if self.tree is not None:
|
|
212
|
+
# Format the XML with proper indentation
|
|
213
|
+
self._indent_xml(self.root)
|
|
214
|
+
return ET.tostring(self.root, encoding="unicode", method="xml")
|
|
215
|
+
msg = "No model loaded. Cannot generate string."
|
|
216
|
+
raise ValueError(msg)
|
|
217
|
+
|
|
218
|
+
def __repr__(self) -> str:
|
|
219
|
+
"""Return the XML string of the model."""
|
|
220
|
+
return self.__str__()
|
|
221
|
+
|
|
222
|
+
def __len__(self) -> int:
|
|
223
|
+
"""Return the number of elements in the model."""
|
|
224
|
+
return len(self.root) if self.root is not None else 0
|
|
225
|
+
|
|
226
|
+
@property
|
|
227
|
+
def xml(self):
|
|
228
|
+
"""XML string of the model."""
|
|
229
|
+
return str(self)
|