From 56b61ae4e7e87b22ea3ce4021367aa5f57a741df Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Fri, 12 Jul 2024 11:22:37 +0200 Subject: [PATCH 1/8] Update MujocoCamera and its usage to be more generic --- src/jaxsim/mujoco/loaders.py | 60 +++++++++++++++++++++++------------- 1 file changed, 39 insertions(+), 21 deletions(-) diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py index ef43a1b4e..42ba97c46 100644 --- a/src/jaxsim/mujoco/loaders.py +++ b/src/jaxsim/mujoco/loaders.py @@ -1,8 +1,10 @@ +from __future__ import annotations + import dataclasses import pathlib import tempfile import warnings -from typing import Any +from typing import Any, Sequence import mujoco as mj import rod.urdf.exporter @@ -161,7 +163,12 @@ def convert( plane_normal: tuple[float, float, float] = (0, 0, 1), heightmap: bool | None = None, heightmap_samples_xy: tuple[int, int] = (101, 101), - cameras: list[dict[str, str]] | dict[str, str] | None = None, + cameras: ( + MujocoCamera + | Sequence[MujocoCamera] + | dict[str, str] + | Sequence[dict[str, str]] + ) = (), ) -> tuple[str, dict[str, Any]]: """ Converts a ROD model to a Mujoco MJCF string. @@ -172,10 +179,10 @@ def convert( plane_normal: The normal vector of the plane. heightmap: Whether to generate a heightmap. heightmap_samples_xy: The number of points in the heightmap grid. - cameras: The list of cameras to add to the scene. + cameras: The custom cameras to add to the scene. Returns: - tuple: A tuple containing the MJCF string and the dictionary of assets. + A tuple containing the MJCF string and the dictionary of assets. """ # ------------------------------------- @@ -478,14 +485,17 @@ def convert( fovy="60", ) - # Add user-defined camera - cameras = cameras if cameras is not None else [] - for camera in cameras if isinstance(cameras, list) else [cameras]: - mj_camera = MujocoCamera.build(**camera) - _ = ET.SubElement( - worldbody_element, "camera", dataclasses.asdict(mj_camera) + # Add user-defined camera. + for camera in cameras if isinstance(cameras, Sequence) else [cameras]: + + mj_camera = ( + camera + if isinstance(camera, MujocoCamera) + else MujocoCamera.build(**camera) ) + _ = ET.SubElement(worldbody_element, "camera", mj_camera.asdict()) + # ------------------------------------------------ # Add a light following the CoM of the first link # ------------------------------------------------ @@ -598,21 +608,29 @@ def convert( @dataclasses.dataclass class MujocoCamera: - name: str - mode: str - pos: str - xyaxes: str - fovy: str + + mode: str = "fixed" + + target: str | None = None + fovy: str = "45" + pos: str = "0 0 0" + + quat: str | None = None + axisangle: str | None = None + xyaxes: str | None = None + zaxis: str | None = None + euler: str | None = None + + name: str | None = None @classmethod - def build(cls, **kwargs): + def build(cls, **kwargs) -> MujocoCamera: + if not all(isinstance(value, str) for value in kwargs.values()): raise ValueError("Values must be strings") - if len(kwargs["pos"].split()) != 3: - raise ValueError("pos must have three values separated by space") + return cls(**kwargs) - if len(kwargs["xyaxes"].split()) != 6: - raise ValueError("xyaxes must have six values separated by space") + def asdict(self) -> dict[str, str]: - return cls(**kwargs) + return {k: v for k, v in dataclasses.asdict(self).items() if v is not None} From b6629ddf42801598d51f16fde2ce639764775e30 Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Fri, 12 Jul 2024 12:09:30 +0200 Subject: [PATCH 2/8] Allow to specify the initial viewpoint of the Mujoco passive viewer --- src/jaxsim/mujoco/visualizer.py | 64 +++++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 2 deletions(-) diff --git a/src/jaxsim/mujoco/visualizer.py b/src/jaxsim/mujoco/visualizer.py index bbf2f400d..96bf8d1cd 100644 --- a/src/jaxsim/mujoco/visualizer.py +++ b/src/jaxsim/mujoco/visualizer.py @@ -1,10 +1,11 @@ import contextlib import pathlib -from typing import ContextManager +from typing import ContextManager, Sequence import mediapy as media import mujoco as mj import mujoco.viewer +import numpy as np import numpy.typing as npt @@ -167,13 +168,72 @@ def open( self, model: mj.MjModel | None = None, data: mj.MjData | None = None, + *, close_on_exit: bool = True, + lookat: Sequence[float | int] | npt.NDArray | None = None, + distance: float | int | npt.NDArray | None = None, + azimut: float | int | npt.NDArray | None = None, + elevation: float | int | npt.NDArray | None = None, ) -> ContextManager[mujoco.viewer.Handle]: - """Context manager to open a viewer.""" + """ + Context manager to open the Mujoco passive viewer. + + Note: + Refer to the Mujoco documentation for details of the camera options: + https://mujoco.readthedocs.io/en/stable/XMLreference.html#visual-global + """ handle = self.open_viewer(model=model, data=data) + handle = MujocoVisualizer.setup_viewer_camera( + viewer=handle, + lookat=lookat, + distance=distance, + azimut=azimut, + elevation=elevation, + ) + try: yield handle finally: _ = handle.close() if close_on_exit else None + + @staticmethod + def setup_viewer_camera( + viewer: mj.viewer.Handle, + *, + lookat: Sequence[float | int] | npt.NDArray | None, + distance: float | int | npt.NDArray | None = None, + azimut: float | int | npt.NDArray | None = None, + elevation: float | int | npt.NDArray | None = None, + ) -> mj.viewer.Handle: + """ + Configure the initial viewpoint of the Mujoco passive viewer. + + Note: + Refer to the Mujoco documentation for details of the camera options: + https://mujoco.readthedocs.io/en/stable/XMLreference.html#visual-global + + Returns: + The viewer with configured camera. + """ + + if lookat is not None: + + lookat_array = np.array(lookat, dtype=float).squeeze() + + if lookat_array.size != 3: + raise ValueError(lookat) + + viewer.cam.lookat = lookat_array + + if distance is not None: + viewer.cam.distance = float(distance) + + if azimut is not None: + viewer.cam.azimuth = float(azimut) % 360 + + if elevation is not None: + viewer.cam.elevation = float(elevation) + + return viewer From 07094fde33229dc1161c53280eeb558341e3690d Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Fri, 12 Jul 2024 11:31:20 +0200 Subject: [PATCH 3/8] Make explicit the camera used by default in the video recorder --- src/jaxsim/mujoco/visualizer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/jaxsim/mujoco/visualizer.py b/src/jaxsim/mujoco/visualizer.py index 96bf8d1cd..70a8f9619 100644 --- a/src/jaxsim/mujoco/visualizer.py +++ b/src/jaxsim/mujoco/visualizer.py @@ -63,18 +63,16 @@ def reset( self.data = data if data is not None else self.data self.model = model if model is not None else self.model - def render_frame(self, camera_name: str | None = None) -> npt.NDArray: + def render_frame(self, camera_name: str = "track") -> npt.NDArray: """Renders a frame.""" - camera_name = camera_name or "track" mujoco.mj_forward(self.model, self.data) self.renderer.update_scene(data=self.data, camera=camera_name) return self.renderer.render() - def record_frame(self, camera_name: str | None = None) -> None: + def record_frame(self, camera_name: str = "track") -> None: """Stores a frame in the buffer.""" - camera_name = camera_name or "track" frame = self.render_frame(camera_name=camera_name) self.frames.append(frame) From 481f4642b298146d069630a86ec35904c7beea9d Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Mon, 15 Jul 2024 12:35:44 +0200 Subject: [PATCH 4/8] Create a custom camera from a more simple target view --- src/jaxsim/mujoco/loaders.py | 80 ++++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py index 42ba97c46..0b876fb84 100644 --- a/src/jaxsim/mujoco/loaders.py +++ b/src/jaxsim/mujoco/loaders.py @@ -7,8 +7,11 @@ from typing import Any, Sequence import mujoco as mj +import numpy as np +import numpy.typing as npt import rod.urdf.exporter from lxml import etree as ET +from scipy.spatial.transform import Rotation def load_rod_model( @@ -631,6 +634,83 @@ def build(cls, **kwargs) -> MujocoCamera: return cls(**kwargs) + @staticmethod + def build_from_target_view( + camera_name: str, + lookat: Sequence[float | int] | npt.NDArray = (0, 0, 0), + distance: float | int | npt.NDArray = 3, + azimut: float | int | npt.NDArray = 90, + elevation: float | int | npt.NDArray = -45, + degrees: bool = True, + **kwargs, + ) -> MujocoCamera: + """ + Create a custom camera that looks at a target point. + + Note: + The choice of the parameters is easier if we imagine to consider a target + frame `T` whose origin is located over the lookat point and having the same + orientation of the world frame `W`. We also introduce a camera frame `C` + whose origin is located over the lower-left corner of the image, and having + the x-axis pointing right and the y-axis pointing up in image coordinates. + The camera renders what it sees in the -z direction of frame `C`. + + Args: + camera_name: The name of the camera. + lookat: The target point to look at (origin of `T`). + distance: + The distance from the target point (displacement between the origins + of `T` and `C`). + azimut: + The rotation around z of the camera. With an angle of 0, the camera + would loot at the target point towards the positive x-axis of `T`. + elevation: + The rotation around the x-axis of the camera frame `C`. Note that if + you want to lift the view angle, the elevation is negative. + degrees: Whether the angles are in degrees or radians. + **kwargs: Additional camera parameters. + + Returns: + The custom camera. + """ + + # Start from a frame whose origin is located over the lookat point. + # We initialize a -90 degrees rotation around the z-axis because due to + # the default camera coordinate system (x pointing right, y pointing up). + W_H_C = np.eye(4) + W_H_C[0:3, 3] = np.array(lookat) + W_H_C[0:3, 0:3] = Rotation.from_euler( + seq="ZX", angles=[-90, 90], degrees=True + ).as_matrix() + + # Process the azimut. + R_az = Rotation.from_euler(seq="Y", angles=azimut, degrees=degrees).as_matrix() + W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_az + + # Process elevation. + R_el = Rotation.from_euler( + seq="X", angles=elevation, degrees=degrees + ).as_matrix() + W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_el + + # Process distance. + tf_distance = np.eye(4) + tf_distance[2, 3] = distance + W_H_C = W_H_C @ tf_distance + + # Extract the position and the quaternion. + p = W_H_C[0:3, 3] + Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat()[np.array([3, 0, 1, 2])] + + return MujocoCamera.build( + name=camera_name, + mode="fixed", + fovy="60", + pos=" ".join(p.astype(str).tolist()), + quat=" ".join(Q.astype(str).tolist()), + **kwargs, + ) + def asdict(self) -> dict[str, str]: return {k: v for k, v in dataclasses.asdict(self).items() if v is not None} From 215878a1920dfe18d313178be6199666ce1cbb96 Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Mon, 15 Jul 2024 10:17:28 +0200 Subject: [PATCH 5/8] Add scipy dependency to [viz] --- environment.yml | 1 + pyproject.toml | 1 + 2 files changed, 2 insertions(+) diff --git a/environment.yml b/environment.yml index 690b36f83..e5937298b 100644 --- a/environment.yml +++ b/environment.yml @@ -30,6 +30,7 @@ dependencies: - lxml - mediapy - mujoco >= 3.0.0 + - scipy >= 1.14.0 # ========================== # Documentation dependencies # ========================== diff --git a/pyproject.toml b/pyproject.toml index e3abd4ba4..c3a3b40ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ viz = [ "lxml", "mediapy", "mujoco >= 3.0.0", + "scipy >= 1.14.0", ] all = [ "jaxsim[style,testing,viz]", From 6b65bc89641620edc47887b33893a9d9b190dd00 Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Mon, 15 Jul 2024 10:17:54 +0200 Subject: [PATCH 6/8] Bump Python version in environment.yml --- environment.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/environment.yml b/environment.yml index e5937298b..3426409d1 100644 --- a/environment.yml +++ b/environment.yml @@ -5,7 +5,7 @@ dependencies: # =========================== # Dependencies from setup.cfg # =========================== - - python=3.11 + - python >= 3.12.0 - coloredlogs - jax >= 0.4.13 - jaxlib >= 0.4.13 From e5ee1dcb958bfff00f785fec6696a52e89502553 Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Mon, 15 Jul 2024 10:22:42 +0200 Subject: [PATCH 7/8] Add fovy to MujocoCamera.build_from_target_view --- src/jaxsim/mujoco/loaders.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py index 0b876fb84..12f884e4d 100644 --- a/src/jaxsim/mujoco/loaders.py +++ b/src/jaxsim/mujoco/loaders.py @@ -641,6 +641,7 @@ def build_from_target_view( distance: float | int | npt.NDArray = 3, azimut: float | int | npt.NDArray = 90, elevation: float | int | npt.NDArray = -45, + fovy: float | int | npt.NDArray = 45, degrees: bool = True, **kwargs, ) -> MujocoCamera: @@ -667,6 +668,7 @@ def build_from_target_view( elevation: The rotation around the x-axis of the camera frame `C`. Note that if you want to lift the view angle, the elevation is negative. + fovy: The field of view of the camera. degrees: Whether the angles are in degrees or radians. **kwargs: Additional camera parameters. @@ -705,7 +707,7 @@ def build_from_target_view( return MujocoCamera.build( name=camera_name, mode="fixed", - fovy="60", + fovy=f"{fovy if degrees else np.rad2deg(fovy)}", pos=" ".join(p.astype(str).tolist()), quat=" ".join(Q.astype(str).tolist()), **kwargs, From b25c33837794a796eda85e1121f6327b379e47c8 Mon Sep 17 00:00:00 2001 From: diegoferigo Date: Mon, 15 Jul 2024 10:43:40 +0200 Subject: [PATCH 8/8] Apply suggestions from code review Co-authored-by: Filippo Luca Ferretti --- src/jaxsim/mujoco/loaders.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py index 12f884e4d..e18faef79 100644 --- a/src/jaxsim/mujoco/loaders.py +++ b/src/jaxsim/mujoco/loaders.py @@ -260,7 +260,6 @@ def convert( parser = ET.XMLParser(remove_blank_text=True) root: ET._Element = ET.fromstring(text=urdf_string.encode(), parser=parser) - import numpy as np # Give a tiny radius to all dummy spheres for geometry in root.findall(".//visual/geometry[sphere]"): @@ -611,6 +610,12 @@ def convert( @dataclasses.dataclass class MujocoCamera: + """ + Helper class storing parameters of a Mujoco camera. + + Refer to the official documentation for more details: + https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-camera + """ mode: str = "fixed" @@ -702,7 +707,7 @@ def build_from_target_view( # Extract the position and the quaternion. p = W_H_C[0:3, 3] - Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat()[np.array([3, 0, 1, 2])] + Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat(scalar_first=True) return MujocoCamera.build( name=camera_name,