jaxsim 0.4.1__py3-none-any.whl → 0.4.1.dev6__py3-none-any.whl
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.
- jaxsim/_version.py +2 -2
- jaxsim/integrators/common.py +33 -24
- jaxsim/mujoco/loaders.py +25 -134
- jaxsim/mujoco/model.py +16 -82
- jaxsim/mujoco/visualizer.py +6 -64
- jaxsim/rbda/utils.py +0 -8
- jaxsim/terrain/terrain.py +1 -1
- {jaxsim-0.4.1.dist-info → jaxsim-0.4.1.dev6.dist-info}/METADATA +1 -2
- {jaxsim-0.4.1.dist-info → jaxsim-0.4.1.dev6.dist-info}/RECORD +12 -12
- {jaxsim-0.4.1.dist-info → jaxsim-0.4.1.dev6.dist-info}/LICENSE +0 -0
- {jaxsim-0.4.1.dist-info → jaxsim-0.4.1.dev6.dist-info}/WHEEL +0 -0
- {jaxsim-0.4.1.dist-info → jaxsim-0.4.1.dev6.dist-info}/top_level.txt +0 -0
jaxsim/_version.py
CHANGED
@@ -12,5 +12,5 @@ __version__: str
|
|
12
12
|
__version_tuple__: VERSION_TUPLE
|
13
13
|
version_tuple: VERSION_TUPLE
|
14
14
|
|
15
|
-
__version__ = version = '0.4.1'
|
16
|
-
__version_tuple__ = version_tuple = (0, 4, 1)
|
15
|
+
__version__ = version = '0.4.1.dev6'
|
16
|
+
__version_tuple__ = version_tuple = (0, 4, 1, 'dev6')
|
jaxsim/integrators/common.py
CHANGED
@@ -5,12 +5,11 @@ from typing import Any, ClassVar, Generic, Protocol, Type, TypeVar
|
|
5
5
|
import jax
|
6
6
|
import jax.numpy as jnp
|
7
7
|
import jax_dataclasses
|
8
|
+
import jaxlie
|
8
9
|
from jax_dataclasses import Static
|
9
10
|
|
10
11
|
import jaxsim.api as js
|
11
|
-
import jaxsim.math
|
12
12
|
import jaxsim.typing as jtp
|
13
|
-
from jaxsim import exceptions
|
14
13
|
from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
|
15
14
|
|
16
15
|
try:
|
@@ -540,38 +539,48 @@ class ExplicitRungeKuttaSO3Mixin:
|
|
540
539
|
`PyTreeType = ODEState` to integrate the quaternion on SO(3).
|
541
540
|
"""
|
542
541
|
|
542
|
+
@classmethod
|
543
|
+
def integrate_rk_stage(
|
544
|
+
cls, x0: js.ode_data.ODEState, t0: Time, dt: TimeStep, k: js.ode_data.ODEState
|
545
|
+
) -> js.ode_data.ODEState:
|
546
|
+
|
547
|
+
op = lambda x0_leaf, k_leaf: x0_leaf + dt * k_leaf
|
548
|
+
xf: js.ode_data.ODEState = jax.tree_util.tree_map(op, x0, k)
|
549
|
+
|
550
|
+
W_Q_B_tf = xf.physics_model.base_quaternion
|
551
|
+
|
552
|
+
return xf.replace(
|
553
|
+
physics_model=xf.physics_model.replace(
|
554
|
+
base_quaternion=W_Q_B_tf / jnp.linalg.norm(W_Q_B_tf)
|
555
|
+
)
|
556
|
+
)
|
557
|
+
|
543
558
|
@classmethod
|
544
559
|
def post_process_state(
|
545
560
|
cls, x0: js.ode_data.ODEState, t0: Time, xf: js.ode_data.ODEState, dt: TimeStep
|
546
561
|
) -> js.ode_data.ODEState:
|
547
562
|
|
548
|
-
#
|
549
|
-
|
563
|
+
# Indices to convert quaternions between serializations.
|
564
|
+
to_xyzw = jnp.array([1, 2, 3, 0])
|
550
565
|
|
551
|
-
#
|
552
|
-
|
553
|
-
|
554
|
-
msg="The SO(3) integrator received a quaternion at t0 that is not unary.",
|
566
|
+
# Get the initial rotation.
|
567
|
+
W_R_B_t0 = jaxlie.SO3.from_quaternion_xyzw(
|
568
|
+
xyzw=x0.physics_model.base_quaternion[to_xyzw]
|
555
569
|
)
|
556
570
|
|
557
|
-
# Get the angular velocity
|
558
|
-
# This
|
559
|
-
#
|
560
|
-
#
|
561
|
-
|
562
|
-
|
563
|
-
|
564
|
-
|
565
|
-
#
|
566
|
-
|
567
|
-
quaternion=W_Q_B_t0,
|
568
|
-
dt=dt,
|
569
|
-
omega=W_ω_WB_t0,
|
570
|
-
omega_in_body_fixed=False,
|
571
|
-
)
|
571
|
+
# Get the final angular velocity.
|
572
|
+
# This is already computed by averaging the kᵢ in RK-based schemes.
|
573
|
+
# Therefore, by using the ω at tf, we obtain a RK scheme operating
|
574
|
+
# on the SO(3) manifold.
|
575
|
+
W_ω_WB_tf = xf.physics_model.base_angular_velocity
|
576
|
+
|
577
|
+
# Integrate the orientation on SO(3).
|
578
|
+
# Note that we left-multiply with the exponential map since the angular
|
579
|
+
# velocity is expressed in the inertial frame.
|
580
|
+
W_R_B_tf = jaxlie.SO3.exp(tangent=dt * W_ω_WB_tf) @ W_R_B_t0
|
572
581
|
|
573
582
|
# Replace the quaternion in the final state.
|
574
583
|
return xf.replace(
|
575
|
-
physics_model=xf.physics_model.replace(base_quaternion=
|
584
|
+
physics_model=xf.physics_model.replace(base_quaternion=W_R_B_tf.wxyz),
|
576
585
|
validate=True,
|
577
586
|
)
|
jaxsim/mujoco/loaders.py
CHANGED
@@ -1,17 +1,12 @@
|
|
1
|
-
from __future__ import annotations
|
2
|
-
|
3
1
|
import dataclasses
|
4
2
|
import pathlib
|
5
3
|
import tempfile
|
6
4
|
import warnings
|
7
|
-
from typing import Any
|
5
|
+
from typing import Any
|
8
6
|
|
9
7
|
import mujoco as mj
|
10
|
-
import numpy as np
|
11
|
-
import numpy.typing as npt
|
12
8
|
import rod.urdf.exporter
|
13
9
|
from lxml import etree as ET
|
14
|
-
from scipy.spatial.transform import Rotation
|
15
10
|
|
16
11
|
|
17
12
|
def load_rod_model(
|
@@ -165,13 +160,7 @@ class RodModelToMjcf:
|
|
165
160
|
considered_joints: list[str] | None = None,
|
166
161
|
plane_normal: tuple[float, float, float] = (0, 0, 1),
|
167
162
|
heightmap: bool | None = None,
|
168
|
-
|
169
|
-
cameras: (
|
170
|
-
MujocoCamera
|
171
|
-
| Sequence[MujocoCamera]
|
172
|
-
| dict[str, str]
|
173
|
-
| Sequence[dict[str, str]]
|
174
|
-
) = (),
|
163
|
+
cameras: list[dict[str, str]] | dict[str, str] | None = None,
|
175
164
|
) -> tuple[str, dict[str, Any]]:
|
176
165
|
"""
|
177
166
|
Converts a ROD model to a Mujoco MJCF string.
|
@@ -181,11 +170,10 @@ class RodModelToMjcf:
|
|
181
170
|
considered_joints: The list of joint names to consider in the conversion.
|
182
171
|
plane_normal: The normal vector of the plane.
|
183
172
|
heightmap: Whether to generate a heightmap.
|
184
|
-
|
185
|
-
cameras: The custom cameras to add to the scene.
|
173
|
+
cameras: The list of cameras to add to the scene.
|
186
174
|
|
187
175
|
Returns:
|
188
|
-
A tuple containing the MJCF string and the dictionary
|
176
|
+
tuple: A tuple containing the MJCF string and the assets dictionary.
|
189
177
|
"""
|
190
178
|
|
191
179
|
# -------------------------------------
|
@@ -260,6 +248,7 @@ class RodModelToMjcf:
|
|
260
248
|
|
261
249
|
parser = ET.XMLParser(remove_blank_text=True)
|
262
250
|
root: ET._Element = ET.fromstring(text=urdf_string.encode(), parser=parser)
|
251
|
+
import numpy as np
|
263
252
|
|
264
253
|
# Give a tiny radius to all dummy spheres
|
265
254
|
for geometry in root.findall(".//visual/geometry[sphere]"):
|
@@ -415,11 +404,9 @@ class RodModelToMjcf:
|
|
415
404
|
asset_element,
|
416
405
|
"hfield",
|
417
406
|
name="terrain",
|
418
|
-
nrow=
|
419
|
-
ncol=
|
420
|
-
|
421
|
-
# when a hfield/heightmap is stored into MjData.
|
422
|
-
size="1 1 1 1",
|
407
|
+
nrow="100",
|
408
|
+
ncol="100",
|
409
|
+
size="5 5 1 1",
|
423
410
|
)
|
424
411
|
if heightmap
|
425
412
|
else None
|
@@ -487,17 +474,14 @@ class RodModelToMjcf:
|
|
487
474
|
fovy="60",
|
488
475
|
)
|
489
476
|
|
490
|
-
# Add user-defined camera
|
491
|
-
|
492
|
-
|
493
|
-
mj_camera = (
|
494
|
-
|
495
|
-
|
496
|
-
else MujocoCamera.build(**camera)
|
477
|
+
# Add user-defined camera
|
478
|
+
cameras = cameras if cameras is not None else []
|
479
|
+
for camera in cameras if isinstance(cameras, list) else [cameras]:
|
480
|
+
mj_camera = MujocoCamera.build(**camera)
|
481
|
+
_ = ET.SubElement(
|
482
|
+
worldbody_element, "camera", dataclasses.asdict(mj_camera)
|
497
483
|
)
|
498
484
|
|
499
|
-
_ = ET.SubElement(worldbody_element, "camera", mj_camera.asdict())
|
500
|
-
|
501
485
|
# ------------------------------------------------
|
502
486
|
# Add a light following the CoM of the first link
|
503
487
|
# ------------------------------------------------
|
@@ -610,114 +594,21 @@ class SdfToMjcf:
|
|
610
594
|
|
611
595
|
@dataclasses.dataclass
|
612
596
|
class MujocoCamera:
|
613
|
-
|
614
|
-
|
615
|
-
|
616
|
-
|
617
|
-
|
618
|
-
"""
|
619
|
-
|
620
|
-
mode: str = "fixed"
|
621
|
-
|
622
|
-
target: str | None = None
|
623
|
-
fovy: str = "45"
|
624
|
-
pos: str = "0 0 0"
|
625
|
-
|
626
|
-
quat: str | None = None
|
627
|
-
axisangle: str | None = None
|
628
|
-
xyaxes: str | None = None
|
629
|
-
zaxis: str | None = None
|
630
|
-
euler: str | None = None
|
631
|
-
|
632
|
-
name: str | None = None
|
597
|
+
name: str
|
598
|
+
mode: str
|
599
|
+
pos: str
|
600
|
+
xyaxes: str
|
601
|
+
fovy: str
|
633
602
|
|
634
603
|
@classmethod
|
635
|
-
def build(cls, **kwargs)
|
636
|
-
|
604
|
+
def build(cls, **kwargs):
|
637
605
|
if not all(isinstance(value, str) for value in kwargs.values()):
|
638
606
|
raise ValueError("Values must be strings")
|
639
607
|
|
640
|
-
|
641
|
-
|
642
|
-
@staticmethod
|
643
|
-
def build_from_target_view(
|
644
|
-
camera_name: str,
|
645
|
-
lookat: Sequence[float | int] | npt.NDArray = (0, 0, 0),
|
646
|
-
distance: float | int | npt.NDArray = 3,
|
647
|
-
azimut: float | int | npt.NDArray = 90,
|
648
|
-
elevation: float | int | npt.NDArray = -45,
|
649
|
-
fovy: float | int | npt.NDArray = 45,
|
650
|
-
degrees: bool = True,
|
651
|
-
**kwargs,
|
652
|
-
) -> MujocoCamera:
|
653
|
-
"""
|
654
|
-
Create a custom camera that looks at a target point.
|
655
|
-
|
656
|
-
Note:
|
657
|
-
The choice of the parameters is easier if we imagine to consider a target
|
658
|
-
frame `T` whose origin is located over the lookat point and having the same
|
659
|
-
orientation of the world frame `W`. We also introduce a camera frame `C`
|
660
|
-
whose origin is located over the lower-left corner of the image, and having
|
661
|
-
the x-axis pointing right and the y-axis pointing up in image coordinates.
|
662
|
-
The camera renders what it sees in the -z direction of frame `C`.
|
663
|
-
|
664
|
-
Args:
|
665
|
-
camera_name: The name of the camera.
|
666
|
-
lookat: The target point to look at (origin of `T`).
|
667
|
-
distance:
|
668
|
-
The distance from the target point (displacement between the origins
|
669
|
-
of `T` and `C`).
|
670
|
-
azimut:
|
671
|
-
The rotation around z of the camera. With an angle of 0, the camera
|
672
|
-
would loot at the target point towards the positive x-axis of `T`.
|
673
|
-
elevation:
|
674
|
-
The rotation around the x-axis of the camera frame `C`. Note that if
|
675
|
-
you want to lift the view angle, the elevation is negative.
|
676
|
-
fovy: The field of view of the camera.
|
677
|
-
degrees: Whether the angles are in degrees or radians.
|
678
|
-
**kwargs: Additional camera parameters.
|
608
|
+
if len(kwargs["pos"].split()) != 3:
|
609
|
+
raise ValueError("pos must have three values separated by space")
|
679
610
|
|
680
|
-
|
681
|
-
|
682
|
-
"""
|
611
|
+
if len(kwargs["xyaxes"].split()) != 6:
|
612
|
+
raise ValueError("xyaxes must have six values separated by space")
|
683
613
|
|
684
|
-
|
685
|
-
# We initialize a -90 degrees rotation around the z-axis because due to
|
686
|
-
# the default camera coordinate system (x pointing right, y pointing up).
|
687
|
-
W_H_C = np.eye(4)
|
688
|
-
W_H_C[0:3, 3] = np.array(lookat)
|
689
|
-
W_H_C[0:3, 0:3] = Rotation.from_euler(
|
690
|
-
seq="ZX", angles=[-90, 90], degrees=True
|
691
|
-
).as_matrix()
|
692
|
-
|
693
|
-
# Process the azimut.
|
694
|
-
R_az = Rotation.from_euler(seq="Y", angles=azimut, degrees=degrees).as_matrix()
|
695
|
-
W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_az
|
696
|
-
|
697
|
-
# Process elevation.
|
698
|
-
R_el = Rotation.from_euler(
|
699
|
-
seq="X", angles=elevation, degrees=degrees
|
700
|
-
).as_matrix()
|
701
|
-
W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_el
|
702
|
-
|
703
|
-
# Process distance.
|
704
|
-
tf_distance = np.eye(4)
|
705
|
-
tf_distance[2, 3] = distance
|
706
|
-
W_H_C = W_H_C @ tf_distance
|
707
|
-
|
708
|
-
# Extract the position and the quaternion.
|
709
|
-
p = W_H_C[0:3, 3]
|
710
|
-
Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat(scalar_first=True)
|
711
|
-
|
712
|
-
return MujocoCamera.build(
|
713
|
-
name=camera_name,
|
714
|
-
mode="fixed",
|
715
|
-
fovy=f"{fovy if degrees else np.rad2deg(fovy)}",
|
716
|
-
pos=" ".join(p.astype(str).tolist()),
|
717
|
-
quat=" ".join(Q.astype(str).tolist()),
|
718
|
-
**kwargs,
|
719
|
-
)
|
720
|
-
|
721
|
-
def asdict(self) -> dict[str, str]:
|
722
|
-
|
723
|
-
return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}
|
614
|
+
return cls(**kwargs)
|
jaxsim/mujoco/model.py
CHANGED
@@ -7,7 +7,6 @@ from typing import Any, Callable
|
|
7
7
|
import mujoco as mj
|
8
8
|
import numpy as np
|
9
9
|
import numpy.typing as npt
|
10
|
-
import xmltodict
|
11
10
|
from scipy.spatial.transform import Rotation
|
12
11
|
|
13
12
|
import jaxsim.typing as jtp
|
@@ -43,27 +42,16 @@ class MujocoModelHelper:
|
|
43
42
|
mjcf_description: str | pathlib.Path,
|
44
43
|
assets: dict[str, Any] | None = None,
|
45
44
|
heightmap: HeightmapCallable | None = None,
|
46
|
-
heightmap_name: str = "terrain",
|
47
|
-
heightmap_radius_xy: tuple[float, float] = (1.0, 1.0),
|
48
45
|
) -> MujocoModelHelper:
|
49
46
|
"""
|
50
|
-
Build a Mujoco model from an
|
47
|
+
Build a Mujoco model from an XML description and an optional assets dictionary.
|
51
48
|
|
52
49
|
Args:
|
53
|
-
mjcf_description:
|
54
|
-
A string containing the XML description of the Mujoco model
|
50
|
+
mjcf_description: A string containing the XML description of the Mujoco model
|
55
51
|
or a path to a file containing the XML description.
|
56
52
|
assets: An optional dictionary containing the assets of the model.
|
57
|
-
heightmap:
|
58
|
-
A function in two variables that returns the height of a terrain
|
53
|
+
heightmap: A function in two variables that returns the height of a terrain
|
59
54
|
in the specified coordinate point.
|
60
|
-
heightmap_name:
|
61
|
-
The default name of the heightmap in the MJCF description
|
62
|
-
to load the corresponding configuration.
|
63
|
-
heightmap_radius_xy:
|
64
|
-
The extension of the heightmap in the x-y surface corresponding to the
|
65
|
-
plane over which the grid of the sampled heightmap is generated.
|
66
|
-
|
67
55
|
Returns:
|
68
56
|
A MujocoModelHelper object.
|
69
57
|
"""
|
@@ -75,61 +63,15 @@ class MujocoModelHelper:
|
|
75
63
|
else mjcf_description
|
76
64
|
)
|
77
65
|
|
78
|
-
|
79
|
-
hfield = None
|
80
|
-
|
81
|
-
else:
|
82
|
-
|
83
|
-
mjcf_description_dict = xmltodict.parse(xml_input=mjcf_description)
|
84
|
-
|
85
|
-
# Create a dictionary of all hfield configurations from the MJCF.
|
86
|
-
hfields = mjcf_description_dict["mujoco"]["asset"].get("hfield", [])
|
87
|
-
hfields = hfields if isinstance(hfields, list) else [hfields]
|
88
|
-
hfields_dict = {hfield["@name"]: hfield for hfield in hfields}
|
89
|
-
|
90
|
-
if heightmap_name not in hfields_dict:
|
91
|
-
raise ValueError(f"Heightmap '{heightmap_name}' not found in MJCF")
|
92
|
-
|
93
|
-
hfield_element = hfields_dict[heightmap_name]
|
94
|
-
|
95
|
-
# Generate the hfield by sampling the heightmap function.
|
96
|
-
hfield = generate_hfield(
|
97
|
-
heightmap=heightmap,
|
98
|
-
samples_xy=(int(hfield_element["@nrow"]), int(hfield_element["@ncol"])),
|
99
|
-
radius_xy=heightmap_radius_xy,
|
100
|
-
)
|
101
|
-
|
102
|
-
# Update dynamically the '/asset/hfield[@name=heightmap_name]@size' attribute
|
103
|
-
# with the information of the sampled points.
|
104
|
-
# This is necessary for correctly rendering the heightmap over the
|
105
|
-
# specified xy area with the correct z elevation.
|
106
|
-
size = [float(el) for el in hfield_element["@size"].split(" ")]
|
107
|
-
size[0], size[1] = heightmap_radius_xy
|
108
|
-
size[2] = 1.0
|
109
|
-
size[3] = max(0, -min(hfield))
|
110
|
-
|
111
|
-
# Replace the 'size' attribute.
|
112
|
-
hfields_dict[heightmap_name]["@size"] = " ".join(str(el) for el in size)
|
113
|
-
|
114
|
-
# Update the hfield elements of the original MJCF.
|
115
|
-
# Only the hfield corresponding to 'heightmap_name' was actually edited.
|
116
|
-
mjcf_description_dict["mujoco"]["asset"]["hfield"] = list(
|
117
|
-
hfields_dict.values()
|
118
|
-
)
|
119
|
-
|
120
|
-
# Serialize the updated MJCF to XML.
|
121
|
-
mjcf_description = xmltodict.unparse(
|
122
|
-
input_dict=mjcf_description_dict, pretty=True
|
123
|
-
)
|
124
|
-
|
125
|
-
# Create the Mujoco model from the XML and, optionally, the dictionary of assets.
|
66
|
+
# Create the Mujoco model from the XML and, optionally, the assets dictionary.
|
126
67
|
model = mj.MjModel.from_xml_string(xml=mjcf_description, assets=assets)
|
127
68
|
data = mj.MjData(model)
|
128
69
|
|
129
|
-
|
130
|
-
|
131
|
-
|
132
|
-
|
70
|
+
if heightmap:
|
71
|
+
nrow = model.hfield_nrow.item()
|
72
|
+
ncol = model.hfield_ncol.item()
|
73
|
+
new_hfield = generate_hfield(heightmap, (nrow, ncol))
|
74
|
+
model.hfield_data = new_hfield
|
133
75
|
|
134
76
|
return MujocoModelHelper(model=model, data=data)
|
135
77
|
|
@@ -443,13 +385,10 @@ class MujocoModelHelper:
|
|
443
385
|
|
444
386
|
|
445
387
|
def generate_hfield(
|
446
|
-
heightmap: HeightmapCallable,
|
447
|
-
samples_xy: tuple[int, int] = (11, 11),
|
448
|
-
radius_xy: tuple[float, float] = (1.0, 1.0),
|
388
|
+
heightmap: HeightmapCallable, size: tuple[int, int] = (10, 10)
|
449
389
|
) -> npt.NDArray:
|
450
390
|
"""
|
451
|
-
|
452
|
-
|
391
|
+
Generates a numpy array representing the heightmap of
|
453
392
|
The map will have the following format:
|
454
393
|
```
|
455
394
|
heightmap[0, 0] heightmap[0, 1] ... heightmap[0, size[1]-1]
|
@@ -459,22 +398,17 @@ def generate_hfield(
|
|
459
398
|
```
|
460
399
|
|
461
400
|
Args:
|
462
|
-
heightmap:
|
463
|
-
A function that takes two arguments (x, y) and returns the height
|
401
|
+
heightmap: A function that takes two arguments (x, y) and returns the height
|
464
402
|
at that point.
|
465
|
-
|
466
|
-
radius_xy:
|
467
|
-
A tuple of two floats representing extension of the heightmap in the
|
468
|
-
x-y surface corresponding to the area over which the grid of the sampled
|
469
|
-
heightmap is generated.
|
403
|
+
size: A tuple of two integers representing the size of the grid.
|
470
404
|
|
471
405
|
Returns:
|
472
|
-
|
406
|
+
np.ndarray: The terrain heightmap
|
473
407
|
"""
|
474
408
|
|
475
409
|
# Generate the grid.
|
476
|
-
x = np.linspace(
|
477
|
-
y = np.linspace(
|
410
|
+
x = np.linspace(0, 1, size[0])
|
411
|
+
y = np.linspace(0, 1, size[1])
|
478
412
|
|
479
413
|
# Generate the heightmap.
|
480
414
|
return np.array([[heightmap(xi, yi) for xi in x] for yi in y]).flatten()
|
jaxsim/mujoco/visualizer.py
CHANGED
@@ -1,11 +1,10 @@
|
|
1
1
|
import contextlib
|
2
2
|
import pathlib
|
3
|
-
from typing import ContextManager
|
3
|
+
from typing import ContextManager
|
4
4
|
|
5
5
|
import mediapy as media
|
6
6
|
import mujoco as mj
|
7
7
|
import mujoco.viewer
|
8
|
-
import numpy as np
|
9
8
|
import numpy.typing as npt
|
10
9
|
|
11
10
|
|
@@ -63,16 +62,18 @@ class MujocoVideoRecorder:
|
|
63
62
|
self.data = data if data is not None else self.data
|
64
63
|
self.model = model if model is not None else self.model
|
65
64
|
|
66
|
-
def render_frame(self, camera_name: str =
|
65
|
+
def render_frame(self, camera_name: str | None = None) -> npt.NDArray:
|
67
66
|
"""Renders a frame."""
|
67
|
+
camera_name = camera_name or "track"
|
68
68
|
|
69
69
|
mujoco.mj_forward(self.model, self.data)
|
70
70
|
self.renderer.update_scene(data=self.data, camera=camera_name)
|
71
71
|
|
72
72
|
return self.renderer.render()
|
73
73
|
|
74
|
-
def record_frame(self, camera_name: str =
|
74
|
+
def record_frame(self, camera_name: str | None = None) -> None:
|
75
75
|
"""Stores a frame in the buffer."""
|
76
|
+
camera_name = camera_name or "track"
|
76
77
|
|
77
78
|
frame = self.render_frame(camera_name=camera_name)
|
78
79
|
self.frames.append(frame)
|
@@ -166,72 +167,13 @@ class MujocoVisualizer:
|
|
166
167
|
self,
|
167
168
|
model: mj.MjModel | None = None,
|
168
169
|
data: mj.MjData | None = None,
|
169
|
-
*,
|
170
170
|
close_on_exit: bool = True,
|
171
|
-
lookat: Sequence[float | int] | npt.NDArray | None = None,
|
172
|
-
distance: float | int | npt.NDArray | None = None,
|
173
|
-
azimut: float | int | npt.NDArray | None = None,
|
174
|
-
elevation: float | int | npt.NDArray | None = None,
|
175
171
|
) -> ContextManager[mujoco.viewer.Handle]:
|
176
|
-
"""
|
177
|
-
Context manager to open the Mujoco passive viewer.
|
178
|
-
|
179
|
-
Note:
|
180
|
-
Refer to the Mujoco documentation for details of the camera options:
|
181
|
-
https://mujoco.readthedocs.io/en/stable/XMLreference.html#visual-global
|
182
|
-
"""
|
172
|
+
"""Context manager to open a viewer."""
|
183
173
|
|
184
174
|
handle = self.open_viewer(model=model, data=data)
|
185
175
|
|
186
|
-
handle = MujocoVisualizer.setup_viewer_camera(
|
187
|
-
viewer=handle,
|
188
|
-
lookat=lookat,
|
189
|
-
distance=distance,
|
190
|
-
azimut=azimut,
|
191
|
-
elevation=elevation,
|
192
|
-
)
|
193
|
-
|
194
176
|
try:
|
195
177
|
yield handle
|
196
178
|
finally:
|
197
179
|
_ = handle.close() if close_on_exit else None
|
198
|
-
|
199
|
-
@staticmethod
|
200
|
-
def setup_viewer_camera(
|
201
|
-
viewer: mj.viewer.Handle,
|
202
|
-
*,
|
203
|
-
lookat: Sequence[float | int] | npt.NDArray | None,
|
204
|
-
distance: float | int | npt.NDArray | None = None,
|
205
|
-
azimut: float | int | npt.NDArray | None = None,
|
206
|
-
elevation: float | int | npt.NDArray | None = None,
|
207
|
-
) -> mj.viewer.Handle:
|
208
|
-
"""
|
209
|
-
Configure the initial viewpoint of the Mujoco passive viewer.
|
210
|
-
|
211
|
-
Note:
|
212
|
-
Refer to the Mujoco documentation for details of the camera options:
|
213
|
-
https://mujoco.readthedocs.io/en/stable/XMLreference.html#visual-global
|
214
|
-
|
215
|
-
Returns:
|
216
|
-
The viewer with configured camera.
|
217
|
-
"""
|
218
|
-
|
219
|
-
if lookat is not None:
|
220
|
-
|
221
|
-
lookat_array = np.array(lookat, dtype=float).squeeze()
|
222
|
-
|
223
|
-
if lookat_array.size != 3:
|
224
|
-
raise ValueError(lookat)
|
225
|
-
|
226
|
-
viewer.cam.lookat = lookat_array
|
227
|
-
|
228
|
-
if distance is not None:
|
229
|
-
viewer.cam.distance = float(distance)
|
230
|
-
|
231
|
-
if azimut is not None:
|
232
|
-
viewer.cam.azimuth = float(azimut) % 360
|
233
|
-
|
234
|
-
if elevation is not None:
|
235
|
-
viewer.cam.elevation = float(elevation)
|
236
|
-
|
237
|
-
return viewer
|
jaxsim/rbda/utils.py
CHANGED
@@ -2,7 +2,6 @@ import jax.numpy as jnp
|
|
2
2
|
|
3
3
|
import jaxsim.api as js
|
4
4
|
import jaxsim.typing as jtp
|
5
|
-
from jaxsim import exceptions
|
6
5
|
from jaxsim.math import StandardGravity
|
7
6
|
|
8
7
|
|
@@ -132,13 +131,6 @@ def process_inputs(
|
|
132
131
|
if W_Q_B.shape != (4,):
|
133
132
|
raise ValueError(W_Q_B.shape, (4,))
|
134
133
|
|
135
|
-
# Check that the quaternion is unary since our RBDAs make this assumption in order
|
136
|
-
# to prevent introducing additional normalizations that would affect AD.
|
137
|
-
exceptions.raise_value_error_if(
|
138
|
-
condition=jnp.logical_not(jnp.allclose(W_Q_B.dot(W_Q_B), 1.0)),
|
139
|
-
msg="A RBDA received a quaternion that is not normalized.",
|
140
|
-
)
|
141
|
-
|
142
134
|
# Pack the 6D base velocity and acceleration.
|
143
135
|
W_v_WB = jnp.hstack([W_vl_WB, W_ω_WB])
|
144
136
|
W_v̇_WB = jnp.hstack([W_v̇l_WB, W_ω̇_WB])
|
jaxsim/terrain/terrain.py
CHANGED
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.1
|
2
2
|
Name: jaxsim
|
3
|
-
Version: 0.4.1
|
3
|
+
Version: 0.4.1.dev6
|
4
4
|
Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
|
5
5
|
Author-email: Diego Ferigo <dgferigo@gmail.com>
|
6
6
|
Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>
|
@@ -82,7 +82,6 @@ Provides-Extra: viz
|
|
82
82
|
Requires-Dist: lxml ; extra == 'viz'
|
83
83
|
Requires-Dist: mediapy ; extra == 'viz'
|
84
84
|
Requires-Dist: mujoco >=3.0.0 ; extra == 'viz'
|
85
|
-
Requires-Dist: scipy >=1.14.0 ; extra == 'viz'
|
86
85
|
|
87
86
|
# JaxSim
|
88
87
|
|
@@ -1,5 +1,5 @@
|
|
1
1
|
jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
|
2
|
-
jaxsim/_version.py,sha256=
|
2
|
+
jaxsim/_version.py,sha256=bHDDEq643CuzU_d4ZEpTqbsQTlTvvqbGfOA9vrbclqI,424
|
3
3
|
jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
|
4
4
|
jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
|
5
5
|
jaxsim/typing.py,sha256=IbFx3UkEXi-cm7UBqMPi58rJAFV_HbZ9E_K4JwfNvVM,753
|
@@ -17,7 +17,7 @@ jaxsim/api/ode.py,sha256=NnLTBvpaT4kXnbjAghXIzLv9DTMJ8bele2iOlUQDv3Q,11028
|
|
17
17
|
jaxsim/api/ode_data.py,sha256=9YZX-SK_KJtoIqG-zYWZsQInb2NA_LtxDn-jtLqm_3U,19759
|
18
18
|
jaxsim/api/references.py,sha256=UA6kSQVBoq-bXSo99EOELf-_MD5MTy2zS0GtG3wQ410,16618
|
19
19
|
jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
|
20
|
-
jaxsim/integrators/common.py,sha256=
|
20
|
+
jaxsim/integrators/common.py,sha256=JXJECMkE-ZSmRt0-5koAw1vnmzkFewZ0aBrnP0bTpZY,20260
|
21
21
|
jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
|
22
22
|
jaxsim/integrators/variable_step.py,sha256=0FCmAZIFnhvQxVbAzNfZgCWN1yMRTGVdBm9UwwaXI1o,21280
|
23
23
|
jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
|
@@ -31,9 +31,9 @@ jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
|
|
31
31
|
jaxsim/math/transform.py,sha256=_5kSnfkS6_vxvjxdw50KeXMjvW8e1OGaumUlk1iGJgc,2969
|
32
32
|
jaxsim/mujoco/__init__.py,sha256=Zo5GAlN1DYKvX8s1hu1j6HntKIbBMLB9Puv9ouaNAZ8,158
|
33
33
|
jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
|
34
|
-
jaxsim/mujoco/loaders.py,sha256=
|
35
|
-
jaxsim/mujoco/model.py,sha256=
|
36
|
-
jaxsim/mujoco/visualizer.py,sha256=
|
34
|
+
jaxsim/mujoco/loaders.py,sha256=FuGPPg9Iq2xLZSqIFzlY7YfPR5uy9rbGlb0cU4gio5A,21000
|
35
|
+
jaxsim/mujoco/model.py,sha256=mObFS77EY97sTFdlFrl69_gf9S_0FHO0W3_hevxQNws,13511
|
36
|
+
jaxsim/mujoco/visualizer.py,sha256=9jfKXkuoaW7Ppo_0m8dogD3SxH13K2strFNDVLtG3hA,5154
|
37
37
|
jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
38
38
|
jaxsim/parsers/kinematic_graph.py,sha256=88d0EmndVJWdcyFJsW25S78Z8F04cUt08RQMyoil1Xw,34734
|
39
39
|
jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
|
@@ -51,18 +51,18 @@ jaxsim/rbda/crba.py,sha256=NhtZO48OUKKor7ddY7mB7h7a6idrmOyf0Vy4p7UCCgI,4724
|
|
51
51
|
jaxsim/rbda/forward_kinematics.py,sha256=OEQYovnLKsWphUKhigmWa_384LwZW3Csp0MKufw4e1M,3415
|
52
52
|
jaxsim/rbda/jacobian.py,sha256=I6mrlkk7Cpq3CE7k_tajOHCbT6vf2pW6vMS0TKNCnng,10725
|
53
53
|
jaxsim/rbda/rnea.py,sha256=UrhcL93fp3pAKlGxOPS6X47L0ferH50bcSMzG55t4zY,7626
|
54
|
-
jaxsim/rbda/utils.py,sha256=
|
54
|
+
jaxsim/rbda/utils.py,sha256=dPRWG8pV7rm3VivP09NE2ttX6wHCRNir4wuixGA_G2Y,5003
|
55
55
|
jaxsim/rbda/contacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
|
56
56
|
jaxsim/rbda/contacts/common.py,sha256=iMKLP30Qft9eGTiHo2iY-UoACJjg1JphA9_pW8wRdjc,2410
|
57
57
|
jaxsim/rbda/contacts/soft.py,sha256=3cDynim_tIgcbzRuqpHN82v4ELlxxK6lR-PG0haSK7Q,15660
|
58
58
|
jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
|
59
|
-
jaxsim/terrain/terrain.py,sha256=
|
59
|
+
jaxsim/terrain/terrain.py,sha256=X4uXkw927nzB6F5OY6WNf4fv85Zc2lk3uP76pjVcMaY,4565
|
60
60
|
jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
|
61
61
|
jaxsim/utils/jaxsim_dataclass.py,sha256=fLl1tY3DDb3lpIhG6BPqA5W34hM84oFzL-5cuz8k-68,11379
|
62
62
|
jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
|
63
63
|
jaxsim/utils/wrappers.py,sha256=GOJQCJc5zwzoEGZB62wnWWGvUUQlXvDxz_A2Q-hFv7c,4027
|
64
|
-
jaxsim-0.4.1.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
|
65
|
-
jaxsim-0.4.1.dist-info/METADATA,sha256=
|
66
|
-
jaxsim-0.4.1.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
|
67
|
-
jaxsim-0.4.1.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
|
68
|
-
jaxsim-0.4.1.dist-info/RECORD,,
|
64
|
+
jaxsim-0.4.1.dev6.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
|
65
|
+
jaxsim-0.4.1.dev6.dist-info/METADATA,sha256=SHBMqYI6b449ER-BGP0ANpyrW84OjD41JgtMbfw41Ss,16778
|
66
|
+
jaxsim-0.4.1.dev6.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
|
67
|
+
jaxsim-0.4.1.dev6.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
|
68
|
+
jaxsim-0.4.1.dev6.dist-info/RECORD,,
|
File without changes
|
File without changes
|
File without changes
|