jaxsim 0.4.0__py3-none-any.whl → 0.4.1__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 +24 -33
- jaxsim/mujoco/loaders.py +134 -25
- jaxsim/mujoco/model.py +82 -16
- jaxsim/mujoco/visualizer.py +64 -6
- jaxsim/rbda/utils.py +8 -0
- jaxsim/terrain/terrain.py +1 -1
- {jaxsim-0.4.0.dist-info → jaxsim-0.4.1.dist-info}/METADATA +2 -1
- {jaxsim-0.4.0.dist-info → jaxsim-0.4.1.dist-info}/RECORD +12 -12
- {jaxsim-0.4.0.dist-info → jaxsim-0.4.1.dist-info}/WHEEL +1 -1
- {jaxsim-0.4.0.dist-info → jaxsim-0.4.1.dist-info}/LICENSE +0 -0
- {jaxsim-0.4.0.dist-info → jaxsim-0.4.1.dist-info}/top_level.txt +0 -0
jaxsim/_version.py
CHANGED
jaxsim/integrators/common.py
CHANGED
@@ -5,11 +5,12 @@ 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
|
9
8
|
from jax_dataclasses import Static
|
10
9
|
|
11
10
|
import jaxsim.api as js
|
11
|
+
import jaxsim.math
|
12
12
|
import jaxsim.typing as jtp
|
13
|
+
from jaxsim import exceptions
|
13
14
|
from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass, Mutability
|
14
15
|
|
15
16
|
try:
|
@@ -539,48 +540,38 @@ class ExplicitRungeKuttaSO3Mixin:
|
|
539
540
|
`PyTreeType = ODEState` to integrate the quaternion on SO(3).
|
540
541
|
"""
|
541
542
|
|
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
|
-
|
558
543
|
@classmethod
|
559
544
|
def post_process_state(
|
560
545
|
cls, x0: js.ode_data.ODEState, t0: Time, xf: js.ode_data.ODEState, dt: TimeStep
|
561
546
|
) -> js.ode_data.ODEState:
|
562
547
|
|
563
|
-
#
|
564
|
-
|
548
|
+
# Extract the initial base quaternion.
|
549
|
+
W_Q_B_t0 = x0.physics_model.base_quaternion
|
565
550
|
|
566
|
-
#
|
567
|
-
|
568
|
-
|
551
|
+
# We assume that the initial quaternion is already unary.
|
552
|
+
exceptions.raise_runtime_error_if(
|
553
|
+
condition=jnp.logical_not(jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0)),
|
554
|
+
msg="The SO(3) integrator received a quaternion at t0 that is not unary.",
|
569
555
|
)
|
570
556
|
|
571
|
-
# Get the
|
572
|
-
# This is
|
573
|
-
# Therefore, by using the ω
|
574
|
-
# on the SO(3) manifold.
|
575
|
-
|
576
|
-
|
577
|
-
|
578
|
-
|
579
|
-
#
|
580
|
-
|
557
|
+
# Get the angular velocity ω to integrate the quaternion.
|
558
|
+
# This velocity ω[t0] is computed in the previous timestep by averaging the kᵢ
|
559
|
+
# corresponding to the active RK-based scheme. Therefore, by using the ω[t0],
|
560
|
+
# we obtain an explicit RK scheme operating on the SO(3) manifold.
|
561
|
+
# Note that the current integrator is not a semi-implicit scheme, therefore
|
562
|
+
# using the final ω[tf] would be not correct.
|
563
|
+
W_ω_WB_t0 = x0.physics_model.base_angular_velocity
|
564
|
+
|
565
|
+
# Integrate the quaternion on SO(3).
|
566
|
+
W_Q_B_tf = jaxsim.math.Quaternion.integration(
|
567
|
+
quaternion=W_Q_B_t0,
|
568
|
+
dt=dt,
|
569
|
+
omega=W_ω_WB_t0,
|
570
|
+
omega_in_body_fixed=False,
|
571
|
+
)
|
581
572
|
|
582
573
|
# Replace the quaternion in the final state.
|
583
574
|
return xf.replace(
|
584
|
-
physics_model=xf.physics_model.replace(base_quaternion=
|
575
|
+
physics_model=xf.physics_model.replace(base_quaternion=W_Q_B_tf),
|
585
576
|
validate=True,
|
586
577
|
)
|
jaxsim/mujoco/loaders.py
CHANGED
@@ -1,12 +1,17 @@
|
|
1
|
+
from __future__ import annotations
|
2
|
+
|
1
3
|
import dataclasses
|
2
4
|
import pathlib
|
3
5
|
import tempfile
|
4
6
|
import warnings
|
5
|
-
from typing import Any
|
7
|
+
from typing import Any, Sequence
|
6
8
|
|
7
9
|
import mujoco as mj
|
10
|
+
import numpy as np
|
11
|
+
import numpy.typing as npt
|
8
12
|
import rod.urdf.exporter
|
9
13
|
from lxml import etree as ET
|
14
|
+
from scipy.spatial.transform import Rotation
|
10
15
|
|
11
16
|
|
12
17
|
def load_rod_model(
|
@@ -160,7 +165,13 @@ class RodModelToMjcf:
|
|
160
165
|
considered_joints: list[str] | None = None,
|
161
166
|
plane_normal: tuple[float, float, float] = (0, 0, 1),
|
162
167
|
heightmap: bool | None = None,
|
163
|
-
|
168
|
+
heightmap_samples_xy: tuple[int, int] = (101, 101),
|
169
|
+
cameras: (
|
170
|
+
MujocoCamera
|
171
|
+
| Sequence[MujocoCamera]
|
172
|
+
| dict[str, str]
|
173
|
+
| Sequence[dict[str, str]]
|
174
|
+
) = (),
|
164
175
|
) -> tuple[str, dict[str, Any]]:
|
165
176
|
"""
|
166
177
|
Converts a ROD model to a Mujoco MJCF string.
|
@@ -170,10 +181,11 @@ class RodModelToMjcf:
|
|
170
181
|
considered_joints: The list of joint names to consider in the conversion.
|
171
182
|
plane_normal: The normal vector of the plane.
|
172
183
|
heightmap: Whether to generate a heightmap.
|
173
|
-
|
184
|
+
heightmap_samples_xy: The number of points in the heightmap grid.
|
185
|
+
cameras: The custom cameras to add to the scene.
|
174
186
|
|
175
187
|
Returns:
|
176
|
-
|
188
|
+
A tuple containing the MJCF string and the dictionary of assets.
|
177
189
|
"""
|
178
190
|
|
179
191
|
# -------------------------------------
|
@@ -248,7 +260,6 @@ class RodModelToMjcf:
|
|
248
260
|
|
249
261
|
parser = ET.XMLParser(remove_blank_text=True)
|
250
262
|
root: ET._Element = ET.fromstring(text=urdf_string.encode(), parser=parser)
|
251
|
-
import numpy as np
|
252
263
|
|
253
264
|
# Give a tiny radius to all dummy spheres
|
254
265
|
for geometry in root.findall(".//visual/geometry[sphere]"):
|
@@ -404,9 +415,11 @@ class RodModelToMjcf:
|
|
404
415
|
asset_element,
|
405
416
|
"hfield",
|
406
417
|
name="terrain",
|
407
|
-
nrow="
|
408
|
-
ncol="
|
409
|
-
size
|
418
|
+
nrow=f"{int(heightmap_samples_xy[0])}",
|
419
|
+
ncol=f"{int(heightmap_samples_xy[1])}",
|
420
|
+
# The following 'size' is a placeholder, it is updated dynamically
|
421
|
+
# when a hfield/heightmap is stored into MjData.
|
422
|
+
size="1 1 1 1",
|
410
423
|
)
|
411
424
|
if heightmap
|
412
425
|
else None
|
@@ -474,14 +487,17 @@ class RodModelToMjcf:
|
|
474
487
|
fovy="60",
|
475
488
|
)
|
476
489
|
|
477
|
-
# Add user-defined camera
|
478
|
-
|
479
|
-
|
480
|
-
mj_camera =
|
481
|
-
|
482
|
-
|
490
|
+
# Add user-defined camera.
|
491
|
+
for camera in cameras if isinstance(cameras, Sequence) else [cameras]:
|
492
|
+
|
493
|
+
mj_camera = (
|
494
|
+
camera
|
495
|
+
if isinstance(camera, MujocoCamera)
|
496
|
+
else MujocoCamera.build(**camera)
|
483
497
|
)
|
484
498
|
|
499
|
+
_ = ET.SubElement(worldbody_element, "camera", mj_camera.asdict())
|
500
|
+
|
485
501
|
# ------------------------------------------------
|
486
502
|
# Add a light following the CoM of the first link
|
487
503
|
# ------------------------------------------------
|
@@ -594,21 +610,114 @@ class SdfToMjcf:
|
|
594
610
|
|
595
611
|
@dataclasses.dataclass
|
596
612
|
class MujocoCamera:
|
597
|
-
|
598
|
-
|
599
|
-
|
600
|
-
|
601
|
-
|
613
|
+
"""
|
614
|
+
Helper class storing parameters of a Mujoco camera.
|
615
|
+
|
616
|
+
Refer to the official documentation for more details:
|
617
|
+
https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-camera
|
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
|
602
633
|
|
603
634
|
@classmethod
|
604
|
-
def build(cls, **kwargs):
|
635
|
+
def build(cls, **kwargs) -> MujocoCamera:
|
636
|
+
|
605
637
|
if not all(isinstance(value, str) for value in kwargs.values()):
|
606
638
|
raise ValueError("Values must be strings")
|
607
639
|
|
608
|
-
|
609
|
-
raise ValueError("pos must have three values separated by space")
|
640
|
+
return cls(**kwargs)
|
610
641
|
|
611
|
-
|
612
|
-
|
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.
|
613
655
|
|
614
|
-
|
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.
|
679
|
+
|
680
|
+
Returns:
|
681
|
+
The custom camera.
|
682
|
+
"""
|
683
|
+
|
684
|
+
# Start from a frame whose origin is located over the lookat point.
|
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}
|
jaxsim/mujoco/model.py
CHANGED
@@ -7,6 +7,7 @@ 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
|
10
11
|
from scipy.spatial.transform import Rotation
|
11
12
|
|
12
13
|
import jaxsim.typing as jtp
|
@@ -42,16 +43,27 @@ class MujocoModelHelper:
|
|
42
43
|
mjcf_description: str | pathlib.Path,
|
43
44
|
assets: dict[str, Any] | None = None,
|
44
45
|
heightmap: HeightmapCallable | None = None,
|
46
|
+
heightmap_name: str = "terrain",
|
47
|
+
heightmap_radius_xy: tuple[float, float] = (1.0, 1.0),
|
45
48
|
) -> MujocoModelHelper:
|
46
49
|
"""
|
47
|
-
Build a Mujoco model from an
|
50
|
+
Build a Mujoco model from an MJCF description.
|
48
51
|
|
49
52
|
Args:
|
50
|
-
mjcf_description:
|
53
|
+
mjcf_description:
|
54
|
+
A string containing the XML description of the Mujoco model
|
51
55
|
or a path to a file containing the XML description.
|
52
56
|
assets: An optional dictionary containing the assets of the model.
|
53
|
-
heightmap:
|
57
|
+
heightmap:
|
58
|
+
A function in two variables that returns the height of a terrain
|
54
59
|
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
|
+
|
55
67
|
Returns:
|
56
68
|
A MujocoModelHelper object.
|
57
69
|
"""
|
@@ -63,15 +75,61 @@ class MujocoModelHelper:
|
|
63
75
|
else mjcf_description
|
64
76
|
)
|
65
77
|
|
66
|
-
|
78
|
+
if heightmap is None:
|
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.
|
67
126
|
model = mj.MjModel.from_xml_string(xml=mjcf_description, assets=assets)
|
68
127
|
data = mj.MjData(model)
|
69
128
|
|
70
|
-
|
71
|
-
|
72
|
-
|
73
|
-
|
74
|
-
model.hfield_data = new_hfield
|
129
|
+
# Store the sampled heightmap into the Mujoco model.
|
130
|
+
if heightmap is not None:
|
131
|
+
assert hfield is not None
|
132
|
+
model.hfield_data = hfield
|
75
133
|
|
76
134
|
return MujocoModelHelper(model=model, data=data)
|
77
135
|
|
@@ -385,10 +443,13 @@ class MujocoModelHelper:
|
|
385
443
|
|
386
444
|
|
387
445
|
def generate_hfield(
|
388
|
-
heightmap: HeightmapCallable,
|
446
|
+
heightmap: HeightmapCallable,
|
447
|
+
samples_xy: tuple[int, int] = (11, 11),
|
448
|
+
radius_xy: tuple[float, float] = (1.0, 1.0),
|
389
449
|
) -> npt.NDArray:
|
390
450
|
"""
|
391
|
-
|
451
|
+
Generate an array with elevation points sampled from a heightmap function.
|
452
|
+
|
392
453
|
The map will have the following format:
|
393
454
|
```
|
394
455
|
heightmap[0, 0] heightmap[0, 1] ... heightmap[0, size[1]-1]
|
@@ -398,17 +459,22 @@ def generate_hfield(
|
|
398
459
|
```
|
399
460
|
|
400
461
|
Args:
|
401
|
-
heightmap:
|
462
|
+
heightmap:
|
463
|
+
A function that takes two arguments (x, y) and returns the height
|
402
464
|
at that point.
|
403
|
-
|
465
|
+
samples_xy: A tuple of two integers representing the size of the grid.
|
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.
|
404
470
|
|
405
471
|
Returns:
|
406
|
-
|
472
|
+
A flat array of the sampled terrain heightmap.
|
407
473
|
"""
|
408
474
|
|
409
475
|
# Generate the grid.
|
410
|
-
x = np.linspace(0,
|
411
|
-
y = np.linspace(
|
476
|
+
x = np.linspace(-radius_xy[0], radius_xy[0], samples_xy[0])
|
477
|
+
y = np.linspace(-radius_xy[1], radius_xy[1], samples_xy[1])
|
412
478
|
|
413
479
|
# Generate the heightmap.
|
414
480
|
return np.array([[heightmap(xi, yi) for xi in x] for yi in y]).flatten()
|
jaxsim/mujoco/visualizer.py
CHANGED
@@ -1,10 +1,11 @@
|
|
1
1
|
import contextlib
|
2
2
|
import pathlib
|
3
|
-
from typing import ContextManager
|
3
|
+
from typing import ContextManager, Sequence
|
4
4
|
|
5
5
|
import mediapy as media
|
6
6
|
import mujoco as mj
|
7
7
|
import mujoco.viewer
|
8
|
+
import numpy as np
|
8
9
|
import numpy.typing as npt
|
9
10
|
|
10
11
|
|
@@ -62,18 +63,16 @@ class MujocoVideoRecorder:
|
|
62
63
|
self.data = data if data is not None else self.data
|
63
64
|
self.model = model if model is not None else self.model
|
64
65
|
|
65
|
-
def render_frame(self, camera_name: str
|
66
|
+
def render_frame(self, camera_name: str = "track") -> npt.NDArray:
|
66
67
|
"""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 = "track") -> None:
|
75
75
|
"""Stores a frame in the buffer."""
|
76
|
-
camera_name = camera_name or "track"
|
77
76
|
|
78
77
|
frame = self.render_frame(camera_name=camera_name)
|
79
78
|
self.frames.append(frame)
|
@@ -167,13 +166,72 @@ class MujocoVisualizer:
|
|
167
166
|
self,
|
168
167
|
model: mj.MjModel | None = None,
|
169
168
|
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,
|
171
175
|
) -> ContextManager[mujoco.viewer.Handle]:
|
172
|
-
"""
|
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
|
+
"""
|
173
183
|
|
174
184
|
handle = self.open_viewer(model=model, data=data)
|
175
185
|
|
186
|
+
handle = MujocoVisualizer.setup_viewer_camera(
|
187
|
+
viewer=handle,
|
188
|
+
lookat=lookat,
|
189
|
+
distance=distance,
|
190
|
+
azimut=azimut,
|
191
|
+
elevation=elevation,
|
192
|
+
)
|
193
|
+
|
176
194
|
try:
|
177
195
|
yield handle
|
178
196
|
finally:
|
179
197
|
_ = 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,6 +2,7 @@ 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
|
5
6
|
from jaxsim.math import StandardGravity
|
6
7
|
|
7
8
|
|
@@ -131,6 +132,13 @@ def process_inputs(
|
|
131
132
|
if W_Q_B.shape != (4,):
|
132
133
|
raise ValueError(W_Q_B.shape, (4,))
|
133
134
|
|
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
|
+
|
134
142
|
# Pack the 6D base velocity and acceleration.
|
135
143
|
W_v_WB = jnp.hstack([W_vl_WB, W_ω_WB])
|
136
144
|
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.
|
3
|
+
Version: 0.4.1
|
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,6 +82,7 @@ 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'
|
85
86
|
|
86
87
|
# JaxSim
|
87
88
|
|
@@ -1,5 +1,5 @@
|
|
1
1
|
jaxsim/__init__.py,sha256=ixsS4dYMPex2wOUUp_rkPnwrPhYzkRh1xO_YuMj3Cr4,2626
|
2
|
-
jaxsim/_version.py,sha256=
|
2
|
+
jaxsim/_version.py,sha256=H6xAc33bJwmO0oYU-O7_wWInmGAifNGi6CO6khT0G1g,411
|
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=iwFykYZxdchqJcmcx8MFWEVijS5Hx9wCNKLKAJdF4gE,20103
|
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=He55jmkC5wQpMhEIDHOXXbqgWNjJ2fx16wOTStp_3PA,25111
|
35
|
+
jaxsim/mujoco/model.py,sha256=EwUPg9BsNv1B7TdDfjZCpC022lDR16AyIAajPJGH7NU,16357
|
36
|
+
jaxsim/mujoco/visualizer.py,sha256=XvMzGSHM-xnOSYl1Vk6bPe6j6ylQmJeLOgxHgL6I1nw,6966
|
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=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
|
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=ctyNANIFSM3tZmamprjaEDcWgUSP0oNJbmT1zw9RjPs,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.
|
65
|
-
jaxsim-0.4.
|
66
|
-
jaxsim-0.4.
|
67
|
-
jaxsim-0.4.
|
68
|
-
jaxsim-0.4.
|
64
|
+
jaxsim-0.4.1.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
|
65
|
+
jaxsim-0.4.1.dist-info/METADATA,sha256=Q5799w-NwG-vzTuZDbKoV25Gr-tO1DTayEBeXTuQZsY,16820
|
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,,
|
File without changes
|
File without changes
|