jaxsim 0.4.1.dev15__py3-none-any.whl → 0.4.1.dev24__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 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.dev15'
16
- __version_tuple__ = version_tuple = (0, 4, 1, 'dev15')
15
+ __version__ = version = '0.4.1.dev24'
16
+ __version_tuple__ = version_tuple = (0, 4, 1, 'dev24')
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(
@@ -161,7 +166,12 @@ class RodModelToMjcf:
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),
164
- cameras: list[dict[str, str]] | dict[str, str] | None = None,
169
+ cameras: (
170
+ MujocoCamera
171
+ | Sequence[MujocoCamera]
172
+ | dict[str, str]
173
+ | Sequence[dict[str, str]]
174
+ ) = (),
165
175
  ) -> tuple[str, dict[str, Any]]:
166
176
  """
167
177
  Converts a ROD model to a Mujoco MJCF string.
@@ -172,10 +182,10 @@ class RodModelToMjcf:
172
182
  plane_normal: The normal vector of the plane.
173
183
  heightmap: Whether to generate a heightmap.
174
184
  heightmap_samples_xy: The number of points in the heightmap grid.
175
- cameras: The list of cameras to add to the scene.
185
+ cameras: The custom cameras to add to the scene.
176
186
 
177
187
  Returns:
178
- tuple: A tuple containing the MJCF string and the dictionary of assets.
188
+ A tuple containing the MJCF string and the dictionary of assets.
179
189
  """
180
190
 
181
191
  # -------------------------------------
@@ -250,7 +260,6 @@ class RodModelToMjcf:
250
260
 
251
261
  parser = ET.XMLParser(remove_blank_text=True)
252
262
  root: ET._Element = ET.fromstring(text=urdf_string.encode(), parser=parser)
253
- import numpy as np
254
263
 
255
264
  # Give a tiny radius to all dummy spheres
256
265
  for geometry in root.findall(".//visual/geometry[sphere]"):
@@ -478,14 +487,17 @@ class RodModelToMjcf:
478
487
  fovy="60",
479
488
  )
480
489
 
481
- # Add user-defined camera
482
- cameras = cameras if cameras is not None else []
483
- for camera in cameras if isinstance(cameras, list) else [cameras]:
484
- mj_camera = MujocoCamera.build(**camera)
485
- _ = ET.SubElement(
486
- worldbody_element, "camera", dataclasses.asdict(mj_camera)
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)
487
497
  )
488
498
 
499
+ _ = ET.SubElement(worldbody_element, "camera", mj_camera.asdict())
500
+
489
501
  # ------------------------------------------------
490
502
  # Add a light following the CoM of the first link
491
503
  # ------------------------------------------------
@@ -598,21 +610,114 @@ class SdfToMjcf:
598
610
 
599
611
  @dataclasses.dataclass
600
612
  class MujocoCamera:
601
- name: str
602
- mode: str
603
- pos: str
604
- xyaxes: str
605
- fovy: str
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
606
633
 
607
634
  @classmethod
608
- def build(cls, **kwargs):
635
+ def build(cls, **kwargs) -> MujocoCamera:
636
+
609
637
  if not all(isinstance(value, str) for value in kwargs.values()):
610
638
  raise ValueError("Values must be strings")
611
639
 
612
- if len(kwargs["pos"].split()) != 3:
613
- raise ValueError("pos must have three values separated by space")
640
+ return cls(**kwargs)
614
641
 
615
- if len(kwargs["xyaxes"].split()) != 6:
616
- raise ValueError("xyaxes must have six values separated by space")
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.
617
655
 
618
- return cls(**kwargs)
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}
@@ -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 | None = None) -> npt.NDArray:
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 | None = None) -> None:
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
- """Context manager to open a viewer."""
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
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: jaxsim
3
- Version: 0.4.1.dev15
3
+ Version: 0.4.1.dev24
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=GGa2wlqOOTBXMs_fQ4BwE7QCLSYjgvCVY3t6lhcCcz4,426
2
+ jaxsim/_version.py,sha256=eeZ5Y19dhLeCQ434QMOmUfyUSYKLUe1iSigAdx-dgIA,426
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
@@ -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=hwT8cjMEZmVaPMs4RzS4UEgC0_c9Od9817UcbNANPBc,21345
34
+ jaxsim/mujoco/loaders.py,sha256=He55jmkC5wQpMhEIDHOXXbqgWNjJ2fx16wOTStp_3PA,25111
35
35
  jaxsim/mujoco/model.py,sha256=EwUPg9BsNv1B7TdDfjZCpC022lDR16AyIAajPJGH7NU,16357
36
- jaxsim/mujoco/visualizer.py,sha256=9jfKXkuoaW7Ppo_0m8dogD3SxH13K2strFNDVLtG3hA,5154
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
@@ -61,8 +61,8 @@ 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.dev15.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
65
- jaxsim-0.4.1.dev15.dist-info/METADATA,sha256=IbvjxjjI0ys4OxT59PyJKmN5Y5xaTDKPT41Q3PTcmbY,16779
66
- jaxsim-0.4.1.dev15.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
67
- jaxsim-0.4.1.dev15.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
68
- jaxsim-0.4.1.dev15.dist-info/RECORD,,
64
+ jaxsim-0.4.1.dev24.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
65
+ jaxsim-0.4.1.dev24.dist-info/METADATA,sha256=dVTGofjfDF-SmXQH3MCUNANxOLYNAI4GnolqdE-MTZw,16826
66
+ jaxsim-0.4.1.dev24.dist-info/WHEEL,sha256=Z4pYXqR_rTB7OWNDYFOm1qRk0RX6GFP2o8LgvP453Hk,91
67
+ jaxsim-0.4.1.dev24.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
68
+ jaxsim-0.4.1.dev24.dist-info/RECORD,,