mujoco-toolbox 0.1.8__tar.gz → 0.2.2__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.1.8 → mujoco_toolbox-0.2.2}/PKG-INFO +29 -82
- mujoco_toolbox-0.2.2/README.md +97 -0
- {mujoco_toolbox-0.1.8 → mujoco_toolbox-0.2.2}/mujoco_toolbox/__init__.py +25 -25
- mujoco_toolbox-0.2.2/mujoco_toolbox/assets.py +29 -0
- mujoco_toolbox-0.2.2/mujoco_toolbox/builder.py +191 -0
- mujoco_toolbox-0.1.8/mujoco_toolbox/Controller.py → mujoco_toolbox-0.2.2/mujoco_toolbox/controller.py +50 -17
- mujoco_toolbox-0.2.2/mujoco_toolbox/utils.py +125 -0
- mujoco_toolbox-0.1.8/mujoco_toolbox/Wrapper.py → mujoco_toolbox-0.2.2/mujoco_toolbox/wrapper.py +289 -193
- mujoco_toolbox-0.2.2/pyproject.toml +127 -0
- mujoco_toolbox-0.1.8/README.md +0 -150
- mujoco_toolbox-0.1.8/mujoco_toolbox/Utils.py +0 -53
- mujoco_toolbox-0.1.8/pyproject.toml +0 -105
- {mujoco_toolbox-0.1.8 → mujoco_toolbox-0.2.2}/LICENSE +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.3
|
|
2
2
|
Name: mujoco-toolbox
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.2
|
|
4
4
|
Summary: A modern package to handle MuJoCo environments
|
|
5
5
|
License: MIT
|
|
6
6
|
Keywords: MuJoCo,robotics,simulation,toolbox,wrapper
|
|
@@ -10,33 +10,35 @@ 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
|
|
13
14
|
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
|
|
14
18
|
Classifier: Topic :: Software Development :: Libraries
|
|
15
|
-
|
|
19
|
+
Requires-Dist: colorama
|
|
16
20
|
Requires-Dist: ffmpeg
|
|
17
21
|
Requires-Dist: ipywidgets
|
|
18
22
|
Requires-Dist: matplotlib
|
|
19
23
|
Requires-Dist: mediapy
|
|
20
24
|
Requires-Dist: mujoco
|
|
21
|
-
Requires-Dist: mypy ; extra == "dev"
|
|
22
25
|
Requires-Dist: numpy
|
|
23
|
-
Requires-Dist: pytest ; extra == "dev"
|
|
24
26
|
Requires-Dist: pyyaml
|
|
25
|
-
Requires-Dist: ruff ; extra == "dev"
|
|
26
27
|
Requires-Dist: scipy
|
|
27
28
|
Requires-Dist: screeninfo
|
|
28
29
|
Requires-Dist: tqdm
|
|
29
30
|
Requires-Dist: trimesh
|
|
30
|
-
Project-URL:
|
|
31
|
-
Project-URL:
|
|
32
|
-
Project-URL:
|
|
33
|
-
Project-URL:
|
|
31
|
+
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
|
|
34
35
|
Description-Content-Type: text/markdown
|
|
35
36
|
|
|
36
37
|
|
|
37
38
|
|
|
38
39
|
|
|
39
40
|
[](https://github.com/MGross21/mujoco-toolbox/actions/workflows/publish.yml)
|
|
41
|
+
[](https://github.com/MGross21/mujoco-toolbox/actions/workflows/docs.yml)
|
|
40
42
|
|
|
41
43
|
# Mujoco Toolbox
|
|
42
44
|
|
|
@@ -48,74 +50,12 @@ Streamlines the MuJoCo Physics Simulator
|
|
|
48
50
|
pip install -U mujoco-toolbox
|
|
49
51
|
```
|
|
50
52
|
|
|
51
|
-
##
|
|
52
|
-
|
|
53
|
-
*Navigate to the Directory that you wish to place the package*
|
|
54
|
-
|
|
55
|
-
<details>
|
|
56
|
-
<summary><b>JUPYTER NOTEBOOK</b></summary>
|
|
57
|
-
</br>
|
|
58
|
-
*Place Commands Inside Cell*
|
|
59
|
-
|
|
60
|
-
One-Time Install:
|
|
61
|
-
|
|
62
|
-
```python
|
|
63
|
-
!git clone https://github.com/MGross21/mujoco-toolbox
|
|
64
|
-
!pip install -U ./mujoco-toolbox
|
|
65
|
-
```
|
|
66
|
-
|
|
67
|
-
Check for Updates:
|
|
68
|
-
|
|
69
|
-
```python
|
|
70
|
-
!git pull origin main
|
|
71
|
-
```
|
|
72
|
-
|
|
73
|
-
Install/Check for Updates:
|
|
74
|
-
|
|
75
|
-
```python
|
|
76
|
-
folder = "mujoco-toolbox"
|
|
77
|
-
repo_url = f"https://github.com/MGross21/{folder}"
|
|
78
|
-
|
|
79
|
-
!if [ -d "{folder}" ]; then cd {folder} && git pull && cd ..; else git clone {repo_url}; fi
|
|
80
|
-
!pip install -U ./{folder}
|
|
81
|
-
```
|
|
82
|
-
|
|
83
|
-
</details>
|
|
84
|
-
</br>
|
|
85
|
-
|
|
86
|
-
<details>
|
|
87
|
-
<summary><b>BASH WINDOW</b></summary>
|
|
88
|
-
</br>
|
|
89
|
-
One-Time Install:
|
|
90
|
-
|
|
91
|
-
```bash
|
|
92
|
-
git clone https://github.com/MGross21/mujoco-toolbox
|
|
93
|
-
pip install -U ./mujoco-toolbox
|
|
94
|
-
```
|
|
95
|
-
|
|
96
|
-
Check for Updates:
|
|
97
|
-
```bash
|
|
98
|
-
git pull origin main
|
|
99
|
-
```
|
|
100
|
-
|
|
101
|
-
Install/Check for Updates:
|
|
53
|
+
## Local Install
|
|
102
54
|
|
|
103
55
|
```bash
|
|
104
|
-
|
|
105
|
-
REPO_URL="https://github.com/MGross21/mujoco-toolbox"
|
|
106
|
-
|
|
107
|
-
if [ -d "$REPO_DIR" ]; then
|
|
108
|
-
cd "$REPO_DIR" && git fetch origin && git pull origin main
|
|
109
|
-
else
|
|
110
|
-
git clone $REPO_URL
|
|
111
|
-
cd "$REPO_DIR"
|
|
112
|
-
fi
|
|
113
|
-
|
|
114
|
-
pip install -U ./mujoco-toolbox
|
|
56
|
+
pip install -U git+https://github.com/MGross21/mujoco-toolbox.git@main
|
|
115
57
|
```
|
|
116
58
|
|
|
117
|
-
</details>
|
|
118
|
-
|
|
119
59
|
## Extra Packages
|
|
120
60
|
|
|
121
61
|
<details>
|
|
@@ -161,26 +101,33 @@ ffmpeg -version
|
|
|
161
101
|
|
|
162
102
|
```python
|
|
163
103
|
import mujoco_toolbox as mjtb
|
|
164
|
-
from mujoco_toolbox import Wrapper
|
|
165
104
|
|
|
166
105
|
# Optional Global Variables
|
|
167
|
-
mjtb.VERBOSITY = True # Additional Performance and Debugging Print Statements (Default=False)
|
|
168
106
|
mjtb.CAPTURE_PARAMETERS = "all" # MjData Simulation Parameters (Default = ['time', 'qpos', 'qvel', 'act', 'qacc', 'xpos', 'xquat', 'xmat', 'ctrl', 'sensordata'])
|
|
169
107
|
|
|
170
108
|
# Example Code
|
|
171
|
-
Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
|
|
109
|
+
mjtb.Wrapper("path/to/your/xml").runSim(render=True).renderMedia()
|
|
172
110
|
```
|
|
173
111
|
|
|
174
112
|
## Pre-Made Controllers
|
|
175
113
|
|
|
176
114
|
```python
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
115
|
+
from mujoco_toolbox import \
|
|
116
|
+
sineController,\
|
|
117
|
+
cosineController,\
|
|
118
|
+
stepController,\
|
|
119
|
+
randomController,\
|
|
120
|
+
realTimeController
|
|
183
121
|
|
|
184
122
|
# Wrapper can use custom controllers as well!
|
|
185
123
|
```
|
|
186
124
|
|
|
125
|
+
## File Support
|
|
126
|
+
|
|
127
|
+
### XML / MJCF (Native)
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
### URDF
|
|
132
|
+
|
|
133
|
+
|
|
@@ -0,0 +1,97 @@
|
|
|
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
|
+
## To Run in Scripts
|
|
65
|
+
|
|
66
|
+
```python
|
|
67
|
+
import mujoco_toolbox as mjtb
|
|
68
|
+
|
|
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()
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
## Pre-Made Controllers
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
from mujoco_toolbox import \
|
|
80
|
+
sineController,\
|
|
81
|
+
cosineController,\
|
|
82
|
+
stepController,\
|
|
83
|
+
randomController,\
|
|
84
|
+
realTimeController
|
|
85
|
+
|
|
86
|
+
# Wrapper can use custom controllers as well!
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
## File Support
|
|
90
|
+
|
|
91
|
+
### XML / MJCF (Native)
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
### URDF
|
|
96
|
+
|
|
97
|
+
|
|
@@ -1,5 +1,4 @@
|
|
|
1
|
-
"""
|
|
2
|
-
Mujoco Toolbox
|
|
1
|
+
"""Mujoco Toolbox.
|
|
3
2
|
==============
|
|
4
3
|
A toolbox for working with MuJoCo simulations. This package provides various utilities and controllers to facilitate
|
|
5
4
|
the simulation process.
|
|
@@ -8,44 +7,47 @@ Modules:
|
|
|
8
7
|
--------
|
|
9
8
|
- Wrapper: Contains the Wrapper class for interfacing with MuJoCo.
|
|
10
9
|
- Controller: Includes sineController, cosineController, and randomController for controlling simulations.
|
|
10
|
+
- Builder: Contains the Builder class for creating and merging MuJoCo models.
|
|
11
11
|
|
|
12
12
|
Constants:
|
|
13
13
|
----------
|
|
14
14
|
- CAPTURE_PARAMETERS: List of MjData fields to capture during simulation.
|
|
15
15
|
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
- __version__ (str): The current version of the package.
|
|
19
|
-
- __author__ (str): The author of the package.
|
|
20
|
-
- __all__ (list): List of public objects of the module, as interpreted by `from module import *`.
|
|
21
|
-
|
|
22
|
-
Notes:
|
|
23
|
-
------
|
|
16
|
+
Notes
|
|
17
|
+
-----
|
|
24
18
|
This package is still under development. Report any issues to https://github.com/MGross21/mujoco-toolbox/issues.
|
|
19
|
+
|
|
25
20
|
"""
|
|
26
21
|
|
|
27
|
-
from .
|
|
28
|
-
|
|
29
|
-
from .
|
|
30
|
-
from .
|
|
22
|
+
from .assets import WORLD_ASSETS, GloveBox
|
|
23
|
+
from .builder import Builder
|
|
24
|
+
from .controller import cosineController, randomController, realTimeController, sineController, stepController
|
|
25
|
+
from .utils import _Platform, print_success, print_warning, timer
|
|
26
|
+
from .wrapper import Wrapper
|
|
31
27
|
|
|
32
|
-
__version__ = "0.
|
|
28
|
+
__version__ = "0.2.2"
|
|
33
29
|
__author__ = "Michael Gross"
|
|
34
30
|
__github_repo__ = "mujoco-toolbox"
|
|
35
31
|
__license__ = "MIT"
|
|
36
32
|
__status__ = "Development"
|
|
37
|
-
|
|
38
|
-
# `from mujoco_toolbox import *` will import these objects
|
|
39
33
|
__all__ = [
|
|
34
|
+
"CAPTURE_PARAMETERS",
|
|
35
|
+
"COMPUTER",
|
|
36
|
+
"WORLD_ASSETS",
|
|
37
|
+
"Builder",
|
|
38
|
+
"GloveBox",
|
|
40
39
|
"Wrapper",
|
|
41
|
-
"sineController",
|
|
42
40
|
"cosineController",
|
|
43
41
|
"randomController",
|
|
42
|
+
"realTimeController",
|
|
43
|
+
"sineController",
|
|
44
44
|
"stepController",
|
|
45
45
|
"timer",
|
|
46
46
|
]
|
|
47
47
|
|
|
48
|
-
|
|
48
|
+
COMPUTER = _Platform() # Singleton Instance
|
|
49
|
+
MAX_GEOM_SCALAR: int = 2 # Scalar value for mujoco.Renderer.max_geom
|
|
50
|
+
CAPTURE_PARAMETERS = [ # MjData default fields to capture during simulation
|
|
49
51
|
"time",
|
|
50
52
|
"qpos",
|
|
51
53
|
"qvel",
|
|
@@ -56,12 +58,10 @@ CAPTURE_PARAMETERS = [
|
|
|
56
58
|
"xmat",
|
|
57
59
|
"ctrl",
|
|
58
60
|
"sensordata",
|
|
59
|
-
]
|
|
60
|
-
|
|
61
|
-
VERBOSITY = False
|
|
61
|
+
]
|
|
62
62
|
|
|
63
|
-
if
|
|
63
|
+
if __version__.split(".")[0] == "0":
|
|
64
64
|
print_warning(
|
|
65
|
-
f"{__package__} ({__version__}) is still under development.",
|
|
66
|
-
f"Report any issues to https://github.com/MGross21/{__github_repo__}/issues
|
|
65
|
+
f"{__package__} (v{__version__}) is still under development.",
|
|
66
|
+
f"Report any issues to https://github.com/MGross21/{__github_repo__}/issues",
|
|
67
67
|
)
|
|
@@ -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,191 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import xml.etree.ElementTree as ET
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class Builder:
|
|
6
|
+
def __init__(self, xml_input: str) -> None:
|
|
7
|
+
"""Initialize with an XML string or a file path. Input is required."""
|
|
8
|
+
self.tree, self.root = self.load_model(xml_input)
|
|
9
|
+
|
|
10
|
+
def load_model(self, xml_input):
|
|
11
|
+
"""Load a MuJoCo model from a file or XML string."""
|
|
12
|
+
if not isinstance(xml_input, str):
|
|
13
|
+
msg = "Input must be either an XML string or a file path"
|
|
14
|
+
raise ValueError(msg)
|
|
15
|
+
|
|
16
|
+
# Check if input is an XML string
|
|
17
|
+
if xml_input.strip().startswith("<"):
|
|
18
|
+
# Try parsing as XML string
|
|
19
|
+
try:
|
|
20
|
+
root = ET.fromstring(xml_input)
|
|
21
|
+
tree = ET.ElementTree(root)
|
|
22
|
+
|
|
23
|
+
# If it's a URDF (robot tag), wrap it in mujoco
|
|
24
|
+
if root.tag == "robot":
|
|
25
|
+
return self.wrap_urdf_in_mujoco(tree)
|
|
26
|
+
|
|
27
|
+
# If root is not mujoco, wrap it
|
|
28
|
+
if root.tag != "mujoco":
|
|
29
|
+
mujoco_elem = ET.Element("mujoco")
|
|
30
|
+
mujoco_elem.append(root)
|
|
31
|
+
tree = ET.ElementTree(mujoco_elem)
|
|
32
|
+
|
|
33
|
+
return tree, tree.getroot()
|
|
34
|
+
except ET.ParseError:
|
|
35
|
+
msg = "Invalid XML string provided"
|
|
36
|
+
raise ValueError(msg)
|
|
37
|
+
else:
|
|
38
|
+
# Treat input as file path
|
|
39
|
+
if not os.path.exists(xml_input):
|
|
40
|
+
msg = f"File not found: {xml_input}"
|
|
41
|
+
raise FileNotFoundError(msg)
|
|
42
|
+
|
|
43
|
+
# Handle URDF files
|
|
44
|
+
if xml_input.lower().endswith(".urdf"):
|
|
45
|
+
return self.wrap_urdf_in_mujoco(ET.parse(xml_input))
|
|
46
|
+
|
|
47
|
+
# Parse regular XML file
|
|
48
|
+
try:
|
|
49
|
+
tree = ET.parse(xml_input)
|
|
50
|
+
root = tree.getroot()
|
|
51
|
+
|
|
52
|
+
# If root is not mujoco, wrap it
|
|
53
|
+
if root.tag != "mujoco":
|
|
54
|
+
mujoco_elem = ET.Element("mujoco")
|
|
55
|
+
mujoco_elem.append(root)
|
|
56
|
+
tree = ET.ElementTree(mujoco_elem)
|
|
57
|
+
|
|
58
|
+
return tree, tree.getroot()
|
|
59
|
+
except ET.ParseError:
|
|
60
|
+
msg = f"Invalid XML file: {xml_input}"
|
|
61
|
+
raise ValueError(msg)
|
|
62
|
+
|
|
63
|
+
def wrap_urdf_in_mujoco(self, urdf_input):
|
|
64
|
+
"""Wrap a URDF inside <mujoco> tags for MuJoCo compatibility.
|
|
65
|
+
Input can be a file path, ElementTree, or XML string.
|
|
66
|
+
"""
|
|
67
|
+
# Handle string input (could be a file path or XML string)
|
|
68
|
+
if isinstance(urdf_input, str):
|
|
69
|
+
# Check if it's an XML string
|
|
70
|
+
if urdf_input.strip().startswith("<"):
|
|
71
|
+
try:
|
|
72
|
+
root = ET.fromstring(urdf_input)
|
|
73
|
+
tree = ET.ElementTree(root)
|
|
74
|
+
except ET.ParseError:
|
|
75
|
+
msg = "Invalid URDF XML string provided"
|
|
76
|
+
raise ValueError(msg)
|
|
77
|
+
else:
|
|
78
|
+
# It's a file path
|
|
79
|
+
try:
|
|
80
|
+
tree = ET.parse(urdf_input)
|
|
81
|
+
except:
|
|
82
|
+
msg = f"Could not parse URDF file: {urdf_input}"
|
|
83
|
+
raise ValueError(msg)
|
|
84
|
+
# Handle ElementTree input
|
|
85
|
+
elif isinstance(urdf_input, ET.ElementTree):
|
|
86
|
+
tree = urdf_input
|
|
87
|
+
else:
|
|
88
|
+
msg = "URDF input must be a string or ElementTree"
|
|
89
|
+
raise TypeError(msg)
|
|
90
|
+
|
|
91
|
+
root = tree.getroot()
|
|
92
|
+
|
|
93
|
+
# Handle the case where the root might already be a mujoco element
|
|
94
|
+
if root.tag == "mujoco":
|
|
95
|
+
return tree, root
|
|
96
|
+
|
|
97
|
+
# Create mujoco element and append the URDF root
|
|
98
|
+
mujoco_elem = ET.Element("mujoco")
|
|
99
|
+
mujoco_elem.append(root)
|
|
100
|
+
|
|
101
|
+
# Create a new tree with the <mujoco> root
|
|
102
|
+
return ET.ElementTree(mujoco_elem), mujoco_elem
|
|
103
|
+
|
|
104
|
+
def merge_tags(self, tag_name, root_1, root_2) -> None:
|
|
105
|
+
"""Merge a specific tag (e.g., worldbody, asset) from two models."""
|
|
106
|
+
section_1 = root_1.find(tag_name)
|
|
107
|
+
section_2 = root_2.find(tag_name)
|
|
108
|
+
|
|
109
|
+
# If target section doesn't exist in model 1 but exists in model 2, create it
|
|
110
|
+
if section_1 is None and section_2 is not None:
|
|
111
|
+
section_1 = ET.SubElement(root_1, tag_name)
|
|
112
|
+
|
|
113
|
+
# Only proceed if both sections exist now
|
|
114
|
+
if section_1 is not None and section_2 is not None:
|
|
115
|
+
# Merge the contents by appending all elements from model 2 to model 1
|
|
116
|
+
for element in list(section_2): # Use list() to avoid modification during iteration
|
|
117
|
+
section_1.append(element)
|
|
118
|
+
|
|
119
|
+
def __add__(self, other):
|
|
120
|
+
"""Implement the + operator for merging two models."""
|
|
121
|
+
model_root = None
|
|
122
|
+
|
|
123
|
+
if isinstance(other, Builder):
|
|
124
|
+
# Merge the current model with another Builder instance
|
|
125
|
+
model_root = other.root
|
|
126
|
+
elif isinstance(other, str):
|
|
127
|
+
# If the input is a string, load as XML
|
|
128
|
+
_, model_root = self.load_model(other)
|
|
129
|
+
else:
|
|
130
|
+
msg = "Addition only supported between Builder objects or with XML strings/files"
|
|
131
|
+
raise TypeError(msg)
|
|
132
|
+
|
|
133
|
+
# List of sections to merge
|
|
134
|
+
sections_to_merge = ["asset", "worldbody", "camera", "light", "contact", "equality",
|
|
135
|
+
"sensor", "actuator", "default", "tendon", "include"]
|
|
136
|
+
|
|
137
|
+
# Merge all relevant sections
|
|
138
|
+
for section in sections_to_merge:
|
|
139
|
+
self.merge_tags(section, self.root, model_root)
|
|
140
|
+
|
|
141
|
+
# Return self to allow for method chaining
|
|
142
|
+
return self
|
|
143
|
+
|
|
144
|
+
def __radd__(self, other):
|
|
145
|
+
"""Implement the reverse + operator for merging two models."""
|
|
146
|
+
if other == 0: # Handle sum() starting value
|
|
147
|
+
return self
|
|
148
|
+
return self.__add__(other)
|
|
149
|
+
|
|
150
|
+
def save(self, file_path) -> None:
|
|
151
|
+
"""Save the merged model to a file."""
|
|
152
|
+
if self.tree is not None:
|
|
153
|
+
# Format the XML with proper indentation before saving
|
|
154
|
+
self._indent_xml(self.root)
|
|
155
|
+
self.tree.write(file_path, encoding="utf-8", xml_declaration=True)
|
|
156
|
+
else:
|
|
157
|
+
msg = "No model loaded. Cannot save."
|
|
158
|
+
raise ValueError(msg)
|
|
159
|
+
|
|
160
|
+
def _indent_xml(self, elem, level=0) -> None:
|
|
161
|
+
"""Add proper indentation to make the XML file more readable."""
|
|
162
|
+
i = "\n" + level * " "
|
|
163
|
+
if len(elem):
|
|
164
|
+
if not elem.text or not elem.text.strip():
|
|
165
|
+
elem.text = i + " "
|
|
166
|
+
if not elem.tail or not elem.tail.strip():
|
|
167
|
+
elem.tail = i
|
|
168
|
+
for sub_elem in elem:
|
|
169
|
+
self._indent_xml(sub_elem, level + 1)
|
|
170
|
+
if not elem[-1].tail or not elem[-1].tail.strip():
|
|
171
|
+
elem[-1].tail = i
|
|
172
|
+
elif level and (not elem.tail or not elem.tail.strip()):
|
|
173
|
+
elem.tail = i
|
|
174
|
+
|
|
175
|
+
def __str__(self) -> str:
|
|
176
|
+
"""Return the XML string of the model."""
|
|
177
|
+
if self.tree is not None:
|
|
178
|
+
# Format the XML with proper indentation
|
|
179
|
+
self._indent_xml(self.root)
|
|
180
|
+
return ET.tostring(self.root, encoding="unicode", method="xml")
|
|
181
|
+
msg = "No model loaded. Cannot generate string."
|
|
182
|
+
raise ValueError(msg)
|
|
183
|
+
|
|
184
|
+
def __repr__(self) -> str:
|
|
185
|
+
"""Return the XML string of the model."""
|
|
186
|
+
return self.__str__()
|
|
187
|
+
|
|
188
|
+
@property
|
|
189
|
+
def xml(self):
|
|
190
|
+
"""Property to quickly grab the XML string of the model."""
|
|
191
|
+
return str(self)
|