jaxsim 0.2.dev188__py3-none-any.whl → 0.6.1.dev2__py3-none-any.whl

jaxsim/mujoco/visualizer.py

@@ -1,15 +1,18 @@
 import contextlib
 import pathlib
-from typing import ContextManager
+from collections.abc import Sequence
 
 import mediapy as media
 import mujoco as mj
 import mujoco.viewer
+import numpy as np
 import numpy.typing as npt
 
 
 class MujocoVideoRecorder:
-    """"""
+    """
+    Video recorder for the MuJoCo passive viewer.
+    """
 
     def __init__(
         self,
@@ -20,7 +23,17 @@ class MujocoVideoRecorder:
         height: int | None = None,
         **kwargs,
     ) -> None:
-        """"""
+        """
+        Initialize the Mujoco video recorder.
+
+        Args:
+            model: The Mujoco model.
+            data: The Mujoco data.
+            fps: The frames per second.
+            width: The width of the video.
+            height: The height of the video.
+            **kwargs: Additional arguments for the renderer.
+        """
 
         width = width if width is not None else model.vis.global_.offwidth
         height = height if height is not None else model.vis.global_.offheight
@@ -45,31 +58,32 @@ class MujocoVideoRecorder:
     def reset(
         self, model: mj.MjModel | None = None, data: mj.MjData | None = None
     ) -> None:
-        """"""
+        """Reset the model and data."""
 
         self.frames = []
 
         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:
-        """"""
-        camera_name = camera_name or "track"
+    def render_frame(self, camera_name: str = "track") -> npt.NDArray:
+        """Render a frame."""
 
         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:
-        """"""
-        camera_name = camera_name or "track"
+    def record_frame(self, camera_name: str = "track") -> None:
+        """Store a frame in the buffer."""
 
         frame = self.render_frame(camera_name=camera_name)
         self.frames.append(frame)
 
     def write_video(self, path: pathlib.Path, exist_ok: bool = False) -> None:
-        """"""
+        """Write the video to a file."""
+
+        # Resolve the path to the video.
+        path = path.expanduser().resolve()
 
         if path.is_dir():
             raise IsADirectoryError(f"The path '{path}' is a directory.")
@@ -77,7 +91,7 @@ class MujocoVideoRecorder:
         if not exist_ok and path.is_file():
             raise FileExistsError(f"The file '{path}' already exists.")
 
-        media.write_video(path=path, images=self.frames, fps=self.fps)
+        media.write_video(path=path, images=np.array(self.frames), fps=self.fps)
 
     @staticmethod
     def compute_down_sampling(original_fps: int, target_min_fps: int) -> int:
@@ -105,12 +119,20 @@ class MujocoVideoRecorder:
 
 
 class MujocoVisualizer:
-    """"""
+    """
+    Visualizer for the MuJoCo passive viewer.
+    """
 
     def __init__(
         self, model: mj.MjModel | None = None, data: mj.MjData | None = None
     ) -> None:
-        """"""
+        """
+        Initialize the Mujoco visualizer.
+
+        Args:
+            model: The Mujoco model.
+            data: The Mujoco data.
+        """
 
         self.data = data
         self.model = model
@@ -121,7 +143,7 @@ class MujocoVisualizer:
         model: mj.MjModel | None = None,
         data: mj.MjData | None = None,
     ) -> None:
-        """"""
+        """Update the viewer with the current model and data."""
 
         data = data if data is not None else self.data
         model = model if model is not None else self.model
@@ -130,15 +152,18 @@ class MujocoVisualizer:
         viewer.sync()
 
     def open_viewer(
-        self, model: mj.MjModel | None = None, data: mj.MjData | None = None
+        self,
+        model: mj.MjModel | None = None,
+        data: mj.MjData | None = None,
+        show_left_ui: bool = False,
     ) -> mj.viewer.Handle:
-        """"""
+        """Open a viewer."""
 
         data = data if data is not None else self.data
         model = model if model is not None else self.model
 
         handle = mj.viewer.launch_passive(
-            model, data, show_left_ui=False, show_right_ui=False
+            model, data, show_left_ui=show_left_ui, show_right_ui=False
         )
 
         return handle
@@ -148,13 +173,73 @@ class MujocoVisualizer:
         self,
         model: mj.MjModel | None = None,
         data: mj.MjData | None = None,
+        *,
+        show_left_ui: bool = False,
         close_on_exit: bool = True,
-    ) -> ContextManager[mujoco.viewer.Handle]:
-        """"""
+        lookat: Sequence[float | int] | npt.NDArray | None = None,
+        distance: float | int | npt.NDArray | None = None,
+        azimuth: float | int | npt.NDArray | None = None,
+        elevation: float | int | npt.NDArray | None = None,
+    ) -> contextlib.AbstractContextManager[mujoco.viewer.Handle]:
+        """
+        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, show_left_ui=show_left_ui)
 
-        handle = self.open_viewer(model=model, data=data)
+        handle = MujocoVisualizer.setup_viewer_camera(
+            viewer=handle,
+            lookat=lookat,
+            distance=distance,
+            azimuth=azimuth,
+            elevation=elevation,
+        )
 
         try:
             yield handle
         finally:
-            handle.close() if close_on_exit else None
+            _ = 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,
+        azimuth: 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 azimuth is not None:
+            viewer.cam.azimuth = float(azimuth) % 360
+
+        if elevation is not None:
+            viewer.cam.elevation = float(elevation)
+
+        return viewer

jaxsim/parsers/__init__.py

@@ -1 +0,0 @@
-from . import descriptions, kinematic_graph

jaxsim/parsers/descriptions/__init__.py

@@ -1,4 +1,10 @@
-from .collision import BoxCollision, CollidablePoint, CollisionShape, SphereCollision
-from .joint import JointDescription, JointDescriptor, JointGenericAxis, JointType
+from .collision import (
+    BoxCollision,
+    CollidablePoint,
+    CollisionShape,
+    MeshCollision,
+    SphereCollision,
+)
+from .joint import JointDescription, JointGenericAxis, JointType
 from .link import LinkDescription
 from .model import ModelDescription

jaxsim/parsers/descriptions/collision.py

@@ -1,11 +1,13 @@
+from __future__ import annotations
+
 import abc
 import dataclasses
-from typing import List
 
 import jax.numpy as jnp
 import numpy as np
 import numpy.typing as npt
 
+import jaxsim.typing as jtp
 from jaxsim import logging
 
 from .link import LinkDescription
@@ -17,10 +19,9 @@ class CollidablePoint:
     Represents a collidable point associated with a parent link.
 
     Attributes:
-        parent_link (LinkDescription): The parent link to which the collidable point is attached.
-        position (npt.NDArray): The position of the collidable point relative to the parent link.
-        enabled (bool): A flag indicating whether the collidable point is enabled for collision detection.
-
+        parent_link: The parent link to which the collidable point is attached.
+        position: The position of the collidable point relative to the parent link.
+        enabled: A flag indicating whether the collidable point is enabled for collision detection.
     """
 
     parent_link: LinkDescription
@@ -29,7 +30,7 @@ class CollidablePoint:
 
     def change_link(
         self, new_link: LinkDescription, new_H_old: npt.NDArray
-    ) -> "CollidablePoint":
+    ) -> CollidablePoint:
         """
         Move the collidable point to a new parent link.
 
@@ -39,8 +40,8 @@ class CollidablePoint:
 
         Returns:
             CollidablePoint: A new collidable point associated with the new parent link.
-
         """
+
         msg = f"Moving collidable point: {self.parent_link.name} -> {new_link.name}"
         logging.debug(msg=msg)
 
@@ -50,7 +51,24 @@ class CollidablePoint:
             enabled=self.enabled,
         )
 
-    def __str__(self):
+    def __hash__(self) -> int:
+
+        return hash(
+            (
+                hash(self.parent_link),
+                hash(tuple(self.position.tolist())),
+                hash(self.enabled),
+            )
+        )
+
+    def __eq__(self, other: CollidablePoint) -> bool:
+
+        if not isinstance(other, CollidablePoint):
+            return False
+
+        return hash(self) == hash(other)
+
+    def __str__(self) -> str:
         return (
             f"{self.__class__.__name__}("
             + f"parent_link={self.parent_link.name}"
@@ -66,11 +84,10 @@ class CollisionShape(abc.ABC):
     Abstract base class for representing collision shapes.
 
     Attributes:
-        collidable_points (List[CollidablePoint]): A list of collidable points associated with the collision shape.
-
+        collidable_points: A list of collidable points associated with the collision shape.
     """
 
-    collidable_points: List[CollidablePoint]
+    collidable_points: tuple[CollidablePoint]
 
     def __str__(self):
         return (
@@ -87,11 +104,25 @@ class BoxCollision(CollisionShape):
     Represents a box-shaped collision shape.
 
     Attributes:
-        center (npt.NDArray): The center of the box in the local frame of the collision shape.
-
+        center: The center of the box in the local frame of the collision shape.
     """
 
-    center: npt.NDArray
+    center: jtp.VectorLike
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                hash(super()),
+                hash(tuple(self.center.tolist())),
+            )
+        )
+
+    def __eq__(self, other: BoxCollision) -> bool:
+
+        if not isinstance(other, BoxCollision):
+            return False
+
+        return hash(self) == hash(other)
 
 
 @dataclasses.dataclass
@@ -100,8 +131,48 @@ class SphereCollision(CollisionShape):
     Represents a spherical collision shape.
 
     Attributes:
-        center (npt.NDArray): The center of the sphere in the local frame of the collision shape.
+        center: The center of the sphere in the local frame of the collision shape.
+    """
+
+    center: jtp.VectorLike
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                hash(super()),
+                hash(tuple(self.center.tolist())),
+            )
+        )
+
+    def __eq__(self, other: BoxCollision) -> bool:
+
+        if not isinstance(other, BoxCollision):
+            return False
+
+        return hash(self) == hash(other)
+
 
+@dataclasses.dataclass
+class MeshCollision(CollisionShape):
+    """
+    Represents a mesh-shaped collision shape.
+
+    Attributes:
+        center: The center of the mesh in the local frame of the collision shape.
     """
 
-    center: npt.NDArray
+    center: jtp.VectorLike
+
+    def __hash__(self) -> int:
+        return hash(
+            (
+                hash(tuple(self.center.tolist())),
+                hash(self.collidable_points),
+            )
+        )
+
+    def __eq__(self, other: MeshCollision) -> bool:
+        if not isinstance(other, MeshCollision):
+            return False
+
+        return hash(self) == hash(other)

jaxsim/parsers/descriptions/joint.py

@@ -1,137 +1,130 @@
+from __future__ import annotations
+
 import dataclasses
-import enum
-from typing import Tuple, Union
+from typing import ClassVar
 
 import jax_dataclasses
 import numpy as np
-import numpy.typing as npt
 
+import jaxsim.typing as jtp
 from jaxsim.utils import JaxsimDataclass, Mutability
 
 from .link import LinkDescription
 
 
-class JointType(enum.IntEnum):
+@dataclasses.dataclass(frozen=True)
+class JointType:
     """
-    Enumeration of joint types for robot joints.
-
-    Args:
-        F: Fixed joint (no movement).
-        R: Revolute joint (rotation).
-        P: Prismatic joint (translation).
-        Rx: Revolute joint with rotation about the X-axis.
-        Ry: Revolute joint with rotation about the Y-axis.
-        Rz: Revolute joint with rotation about the Z-axis.
-        Px: Prismatic joint with translation along the X-axis.
-        Py: Prismatic joint with translation along the Y-axis.
-        Pz: Prismatic joint with translation along the Z-axis.
+    Enumeration of joint types.
     """
 
-    F = enum.auto()  # Fixed
-    R = enum.auto()  # Revolute
-    P = enum.auto()  # Prismatic
-
-    # Revolute joints, single axis
-    Rx = enum.auto()
-    Ry = enum.auto()
-    Rz = enum.auto()
-
-    # Prismatic joints, single axis
-    Px = enum.auto()
-    Py = enum.auto()
-    Pz = enum.auto()
+    Fixed: ClassVar[int] = 0
+    Revolute: ClassVar[int] = 1
+    Prismatic: ClassVar[int] = 2
 
 
-@dataclasses.dataclass
-class JointDescriptor:
+@jax_dataclasses.pytree_dataclass
+class JointGenericAxis:
     """
-    Description of a joint type with a specific code.
-
-    Args:
-        code (JointType): The code representing the joint type.
-
+    A joint requiring the specification of a 3D axis.
     """
 
-    code: JointType
+    # The axis of rotation or translation of the joint (must have norm 1).
+    axis: jtp.Vector
 
     def __hash__(self) -> int:
-        return hash(self.__repr__())
-
-
-@dataclasses.dataclass
-class JointGenericAxis(JointDescriptor):
-    """
-    Description of a joint type with a generic axis.
-
-    Attributes:
-        axis (npt.NDArray): The axis of rotation or translation for the joint.
 
-    """
-
-    axis: npt.NDArray
+        return hash(tuple(self.axis.tolist()))
 
-    def __post_init__(self):
-        if np.allclose(self.axis, 0.0):
-            raise ValueError(self.axis)
+    def __eq__(self, other: JointGenericAxis) -> bool:
 
-    def __eq__(self, other):
-        return super().__eq__(other) and np.allclose(self.axis, other.axis)
+        if not isinstance(other, JointGenericAxis):
+            return False
 
-    def __hash__(self) -> int:
-        return hash(self.__repr__())
+        return hash(self) == hash(other)
 
 
-@jax_dataclasses.pytree_dataclass
+@jax_dataclasses.pytree_dataclass(eq=False, unsafe_hash=False)
 class JointDescription(JaxsimDataclass):
     """
     In-memory description of a robot link.
 
     Attributes:
-        name (str): The name of the joint.
-        axis (npt.NDArray): The axis of rotation or translation for the joint.
-        pose (npt.NDArray): The pose transformation matrix of the joint.
-        jtype (Union[JointType, JointDescriptor]): The type of the joint.
-        child (LinkDescription): The child link attached to the joint.
-        parent (LinkDescription): The parent link attached to the joint.
-        index (Optional[int]): An optional index for the joint.
-        friction_static (float): The static friction coefficient for the joint.
-        friction_viscous (float): The viscous friction coefficient for the joint.
-        position_limit_damper (float): The damper coefficient for position limits.
-        position_limit_spring (float): The spring coefficient for position limits.
-        position_limit (Tuple[float, float]): The position limits for the joint.
-        initial_position (Union[float, npt.NDArray]): The initial position of the joint.
-
+        name: The name of the joint.
+        axis: The axis of rotation or translation for the joint.
+        pose: The pose transformation matrix of the joint.
+        jtype: The type of the joint.
+        child: The child link attached to the joint.
+        parent: The parent link attached to the joint.
+        index: An optional index for the joint.
+        friction_static: The static friction coefficient for the joint.
+        friction_viscous: The viscous friction coefficient for the joint.
+        position_limit_damper: The damper coefficient for position limits.
+        position_limit_spring: The spring coefficient for position limits.
+        position_limit: The position limits for the joint.
+        initial_position: The initial position of the joint.
     """
 
     name: jax_dataclasses.Static[str]
-    axis: npt.NDArray
-    pose: npt.NDArray
-    jtype: jax_dataclasses.Static[Union[JointType, JointDescriptor]]
+    axis: jtp.Vector
+    pose: jtp.Matrix
+    jtype: jax_dataclasses.Static[jtp.IntLike]
     child: LinkDescription = dataclasses.dataclass(repr=False)
     parent: LinkDescription = dataclasses.dataclass(repr=False)
 
-    index: int | None = None
+    index: jtp.IntLike | None = None
+
+    friction_static: jtp.FloatLike = 0.0
+    friction_viscous: jtp.FloatLike = 0.0
 
-    friction_static: float = 0.0
-    friction_viscous: float = 0.0
+    position_limit_damper: jtp.FloatLike = 0.0
+    position_limit_spring: jtp.FloatLike = 0.0
 
-    position_limit_damper: float = 0.0
-    position_limit_spring: float = 0.0
+    position_limit: tuple[jtp.FloatLike, jtp.FloatLike] = (0.0, 0.0)
+    initial_position: jtp.FloatLike | jtp.VectorLike = 0.0
 
-    position_limit: Tuple[float, float] = (0.0, 0.0)
-    initial_position: Union[float, npt.NDArray] = 0.0
+    motor_inertia: jtp.FloatLike = 0.0
+    motor_viscous_friction: jtp.FloatLike = 0.0
+    motor_gear_ratio: jtp.FloatLike = 1.0
 
-    motor_inertia: float = 0.0
-    motor_viscous_friction: float = 0.0
-    motor_gear_ratio: float = 1.0
+    def __post_init__(self) -> None:
 
-    def __post_init__(self):
         if self.axis is not None:
+
             with self.mutable_context(
                 mutability=Mutability.MUTABLE, restore_after_exception=False
             ):
                 norm_of_axis = np.linalg.norm(self.axis)
                 self.axis = self.axis / norm_of_axis
 
+    def __eq__(self, other: JointDescription) -> bool:
+
+        if not isinstance(other, JointDescription):
+            return False
+
+        return hash(self) == hash(other)
+
     def __hash__(self) -> int:
-        return hash(self.__repr__())
+
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return hash(
+            (
+                hash(self.name),
+                HashedNumpyArray.hash_of_array(self.axis),
+                HashedNumpyArray.hash_of_array(self.pose),
+                hash(int(self.jtype)),
+                hash(self.child),
+                hash(self.parent),
+                hash(int(self.index)) if self.index is not None else 0,
+                HashedNumpyArray.hash_of_array(self.friction_static),
+                HashedNumpyArray.hash_of_array(self.friction_viscous),
+                HashedNumpyArray.hash_of_array(self.position_limit_damper),
+                HashedNumpyArray.hash_of_array(self.position_limit_spring),
+                HashedNumpyArray.hash_of_array(self.position_limit),
+                HashedNumpyArray.hash_of_array(self.initial_position),
+                HashedNumpyArray.hash_of_array(self.motor_inertia),
+                HashedNumpyArray.hash_of_array(self.motor_viscous_friction),
+                HashedNumpyArray.hash_of_array(self.motor_gear_ratio),
+            ),
+        )