jaxsim 0.3.1.dev62__py3-none-any.whl → 0.3.1.dev94__py3-none-any.whl

jaxsim/parsers/rod/utils.py

@@ -1,12 +1,11 @@
 import os
 
-import jaxlie
 import numpy as np
 import numpy.typing as npt
 import rod
 
 import jaxsim.typing as jtp
-from jaxsim.math import Inertia
+from jaxsim.math import Adjoint, Inertia
 from jaxsim.parsers import descriptions
 
 
@@ -21,10 +20,10 @@ def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix:
         The 6D inertia matrix of the link expressed in the link frame.
     """
 
-    # Extract the "mass" element
+    # Extract the "mass" element.
     m = inertial.mass
 
-    # Extract the "inertia" element
+    # Extract the "inertia" element.
     inertia_element = inertial.inertia
 
     ixx = inertia_element.ixx
@@ -34,7 +33,7 @@ def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix:
     ixz = inertia_element.ixz if inertia_element.ixz is not None else 0.0
     iyz = inertia_element.iyz if inertia_element.iyz is not None else 0.0
 
-    # Build the 3x3 inertia matrix expressed in the CoM
+    # Build the 3x3 inertia matrix expressed in the CoM.
     I_CoM = np.array(
         [
             [ixx, ixy, ixz],
@@ -43,17 +42,16 @@ def from_sdf_inertial(inertial: rod.Inertial) -> jtp.Matrix:
         ]
     )
 
-    # Build the 6x6 generalized inertia at the CoM
+    # Build the 6x6 generalized inertia at the CoM.
     M_CoM = Inertia.to_sixd(mass=m, com=np.zeros(3), I=I_CoM)
 
-    # Compute the transform from the inertial frame (CoM) to the link frame
+    # Compute the transform from the inertial frame (CoM) to the link frame.
     L_H_CoM = inertial.pose.transform() if inertial.pose is not None else np.eye(4)
 
-    # We need its inverse
-    CoM_H_L = jaxlie.SE3.from_matrix(matrix=L_H_CoM).inverse()
-    CoM_X_L = CoM_H_L.adjoint()
+    # We need its inverse.
+    CoM_X_L = Adjoint.from_transform(transform=L_H_CoM, inverse=True)
 
-    # Express the CoM inertia matrix in the link frame L
+    # Express the CoM inertia matrix in the link frame L.
     M_L = CoM_X_L.T @ M_CoM @ CoM_X_L
 
     return M_L.astype(dtype=float)

jaxsim/rbda/aba.py

@@ -102,7 +102,7 @@ def aba(
     i_X_0 = jnp.zeros(shape=(model.number_of_links(), 6, 6))
     i_X_0 = i_X_0.at[0].set(jnp.eye(6))
 
-    # Initialize base quantities
+    # Initialize base quantities.
     if model.floating_base():
 
         # Base velocity v₀ in body-fixed representation.
@@ -121,10 +121,7 @@ def aba(
     # Pass 1
     # ======
 
-    Pass1Carry = tuple[
-        jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
-    ]
-
+    Pass1Carry = tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix]
     pass_1_carry: Pass1Carry = (v, c, MA, pA, i_X_0)
 
     # Propagate kinematics and initialize AB inertia and AB bias forces.
@@ -178,10 +175,7 @@ def aba(
     d = jnp.zeros(shape=(model.number_of_links(), 1))
     u = jnp.zeros(shape=(model.number_of_links(), 1))
 
-    Pass2Carry = tuple[
-        jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
-    ]
-
+    Pass2Carry = tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix]
     pass_2_carry: Pass2Carry = (U, d, u, MA, pA)
 
     def loop_body_pass2(carry: Pass2Carry, i: jtp.Int) -> tuple[Pass2Carry, None]:
@@ -204,8 +198,8 @@ def aba(
 
         # Propagate them to the parent, handling the base link.
         def propagate(
-            MA_pA: tuple[jtp.MatrixJax, jtp.MatrixJax]
-        ) -> tuple[jtp.MatrixJax, jtp.MatrixJax]:
+            MA_pA: tuple[jtp.Matrix, jtp.Matrix]
+        ) -> tuple[jtp.Matrix, jtp.Matrix]:
 
             MA, pA = MA_pA
 
@@ -248,7 +242,7 @@ def aba(
     s̈ = jnp.zeros_like(s)
     a = jnp.zeros_like(v).at[0].set(a0)
 
-    Pass3Carry = tuple[jtp.MatrixJax, jtp.VectorJax]
+    Pass3Carry = tuple[jtp.Matrix, jtp.Vector]
     pass_3_carry = (a, s̈)
 
     def loop_body_pass3(carry: Pass3Carry, i: jtp.Int) -> tuple[Pass3Carry, None]:

jaxsim/rbda/collidable_points.py

@@ -80,7 +80,7 @@ def collidable_points_pos_vel(
     # Propagate kinematics
     # ====================
 
-    PropagateTransformsCarry = tuple[jtp.MatrixJax, jtp.Matrix]
+    PropagateTransformsCarry = tuple[jtp.Matrix, jtp.Matrix]
     propagate_transforms_carry: PropagateTransformsCarry = (W_X_i, W_v_Wi)
 
     def propagate_kinematics(
@@ -97,7 +97,7 @@ def collidable_points_pos_vel(
         W_Xi_i = W_X_i[λ[i]] @ λi_X_i
         W_X_i = W_X_i.at[i].set(W_Xi_i)
 
-        # Propagate the 6D velocity
+        # Propagate the 6D velocity.
         W_vi_Wi = W_v_Wi[λ[i]] + W_X_i[i] @ (S[i] * ṡ[ii]).squeeze()
         W_v_Wi = W_v_Wi.at[i].set(W_vi_Wi)
 
@@ -118,14 +118,15 @@ def collidable_points_pos_vel(
     # ==================================================
 
     def process_point_kinematics(
-        Li_p_C: jtp.VectorJax, parent_body: jtp.Int
-    ) -> tuple[jtp.VectorJax, jtp.VectorJax]:
-        # Compute the position of the collidable point
+        Li_p_C: jtp.Vector, parent_body: jtp.Int
+    ) -> tuple[jtp.Vector, jtp.Vector]:
+
+        # Compute the position of the collidable point.
         W_p_Ci = (
             Adjoint.to_transform(adjoint=W_X_i[parent_body]) @ jnp.hstack([Li_p_C, 1])
         )[0:3]
 
-        # Compute the linear part of the mixed velocity Ci[W]_v_{W,Ci}
+        # Compute the linear part of the mixed velocity Ci[W]_v_{W,Ci}.
         CW_vl_WCi = (
             jnp.block([jnp.eye(3), -Skew.wedge(vector=W_p_Ci).squeeze()])
             @ W_v_Wi[parent_body].squeeze()
@@ -133,7 +134,7 @@ def collidable_points_pos_vel(
 
         return W_p_Ci, CW_vl_WCi
 
-    # Process all the collidable points in parallel
+    # Process all the collidable points in parallel.
     W_p_Ci, CW_vl_WC = jax.vmap(process_point_kinematics)(
         model.kin_dyn_parameters.contact_parameters.point,
         jnp.array(model.kin_dyn_parameters.contact_parameters.body),

jaxsim/rbda/contacts/soft.py

@@ -105,24 +105,24 @@ class SoftContactsParams(ContactsParams):
             - ξ < 1.0: under-damped
         """
 
-        # Use symbols for input parameters
+        # Use symbols for input parameters.
         ξ = damping_ratio
         δ_max = max_penetration
         μc = static_friction_coefficient
 
-        # Compute the total mass of the model
+        # Compute the total mass of the model.
         m = jnp.array(model.kin_dyn_parameters.link_parameters.mass).sum()
 
-        # Rename the standard gravity
+        # Rename the standard gravity.
         g = standard_gravity
 
-        # Compute the average support force on each collidable point
+        # Compute the average support force on each collidable point.
         f_average = m * g / number_of_active_collidable_points_steady_state
 
-        # Compute the stiffness to get the desired steady-state penetration
+        # Compute the stiffness to get the desired steady-state penetration.
         K = f_average / jnp.power(δ_max, 3 / 2)
 
-        # Compute the damping using the damping ratio
+        # Compute the damping using the damping ratio.
         critical_damping = 2 * jnp.sqrt(K * m)
         D = ξ * critical_damping
 
@@ -151,14 +151,16 @@ class SoftContacts(ContactModel):
         default_factory=SoftContactsParams
     )
 
-    terrain: Terrain = dataclasses.field(default_factory=FlatTerrain)
+    terrain: jax_dataclasses.Static[Terrain] = dataclasses.field(
+        default_factory=FlatTerrain
+    )
 
     def compute_contact_forces(
         self,
         position: jtp.Vector,
         velocity: jtp.Vector,
         tangential_deformation: jtp.Vector,
-    ) -> tuple[jtp.Vector, tuple[jtp.Vector, None]]:
+    ) -> tuple[jtp.Vector, tuple[jtp.Vector]]:
         """
         Compute the contact forces and material deformation rate.
 
@@ -188,18 +190,18 @@ class SoftContacts(ContactModel):
         # Normal force computation
         # ========================
 
-        # Unpack the position of the collidable point
+        # Unpack the position of the collidable point.
         px, py, pz = W_p_C = position.squeeze()
         vx, vy, vz = W_ṗ_C = velocity.squeeze()
 
-        # Compute the terrain normal and the contact depth
+        # Compute the terrain normal and the contact depth.
         n̂ = self.terrain.normal(x=px, y=py).squeeze()
         h = jnp.array([0, 0, self.terrain.height(x=px, y=py) - pz])
 
-        # Compute the penetration depth normal to the terrain
+        # Compute the penetration depth normal to the terrain.
         δ = jnp.maximum(0.0, jnp.dot(h, n̂))
 
-        # Compute the penetration normal velocity
+        # Compute the penetration normal velocity.
         δ̇ = -jnp.dot(W_ṗ_C, n̂)
 
         # Non-linear spring-damper model.
@@ -210,10 +212,10 @@ class SoftContacts(ContactModel):
             on_false=jnp.array(0.0),
         )
 
-        # Prevent negative normal forces that might occur when δ̇ is largely negative
+        # Prevent negative normal forces that might occur when δ̇ is largely negative.
         force_normal_mag = jnp.maximum(0.0, force_normal_mag)
 
-        # Compute the 3D linear force in C[W] frame
+        # Compute the 3D linear force in C[W] frame.
         force_normal = force_normal_mag * n̂
 
         # ====================================
@@ -230,11 +232,11 @@ class SoftContacts(ContactModel):
         )
 
         def with_no_friction():
-            # Compute 6D mixed force in C[W]
+            # Compute 6D mixed force in C[W].
             CW_f_lin = force_normal
             CW_f = jnp.hstack([force_normal, jnp.zeros_like(CW_f_lin)])
 
-            # Compute lin-ang 6D forces (inertial representation)
+            # Compute lin-ang 6D forces (inertial representation).
             W_f = W_Xf_CW @ CW_f
 
             return W_f, (ṁ,)
@@ -258,32 +260,32 @@ class SoftContacts(ContactModel):
                 return jnp.zeros(6), (ṁ,)
 
             def below_terrain():
-                # Decompose the velocity in normal and tangential components
+                # Decompose the velocity in normal and tangential components.
                 v_normal = jnp.dot(W_ṗ_C, n̂) * n̂
                 v_tangential = W_ṗ_C - v_normal
 
-                # Compute the tangential force. If inside the friction cone, the contact
+                # Compute the tangential force. If inside the friction cone, the contact.
                 f_tangential = -jnp.sqrt(δ + 1e-12) * (K * m + D * v_tangential)
 
                 def sticking_contact():
-                    # Sum the normal and tangential forces, and create the 6D force
+                    # Sum the normal and tangential forces, and create the 6D force.
                     CW_f_stick = force_normal + f_tangential
                     CW_f = jnp.hstack([CW_f_stick, jnp.zeros(3)])
 
-                    # In this case the 3D material deformation is the tangential velocity
+                    # In this case the 3D material deformation is the tangential velocity.
                     ṁ = v_tangential
 
                     # Return the 6D force in the contact frame and
-                    # the deformation derivative
+                    # the deformation derivative.
                     return CW_f, ṁ
 
                 def slipping_contact():
-                    # Project the force to the friction cone boundary
+                    # Project the force to the friction cone boundary.
                     f_tangential_projected = (μ * force_normal_mag) * (
                         f_tangential / jnp.maximum(jnp.linalg.norm(f_tangential), 1e-9)
                     )
 
-                    # Sum the normal and tangential forces, and create the 6D force
+                    # Sum the normal and tangential forces, and create the 6D force.
                     CW_f_slip = force_normal + f_tangential_projected
                     CW_f = jnp.hstack([CW_f_slip, jnp.zeros(3)])
 
@@ -297,7 +299,7 @@ class SoftContacts(ContactModel):
                     ṁ = (f_tangential_projected - α * m) / β
 
                     # Return the 6D force in the contact frame and
-                    # the deformation derivative
+                    # the deformation derivative.
                     return CW_f, ṁ
 
                 CW_f, ṁ = jax.lax.cond(
@@ -307,10 +309,10 @@ class SoftContacts(ContactModel):
                     operand=None,
                 )
 
-                # Express the 6D force in the world frame
+                # Express the 6D force in the world frame.
                 W_f = W_Xf_CW @ CW_f
 
-                # Return the 6D force in the world frame and the deformation derivative
+                # Return the 6D force in the world frame and the deformation derivative.
                 return W_f, (ṁ,)
 
             # (W_f, (ṁ,))
@@ -321,7 +323,7 @@ class SoftContacts(ContactModel):
                 operand=None,
             )
 
-        # (W_f, ṁ)
+        # (W_f, (ṁ,))
         return jax.lax.cond(
             pred=(μ == 0.0),
             true_fun=lambda _: with_no_friction(),

jaxsim/rbda/crba.py

@@ -45,7 +45,7 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
     # Propagate kinematics
     # ====================
 
-    ForwardPassCarry = tuple[jtp.MatrixJax]
+    ForwardPassCarry = tuple[jtp.Matrix]
     forward_pass_carry: ForwardPassCarry = (i_X_0,)
 
     def propagate_kinematics(
@@ -71,7 +71,7 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
     M = jnp.zeros(shape=(6 + model.dofs(), 6 + model.dofs()))
 
-    BackwardPassCarry = tuple[jtp.MatrixJax, jtp.MatrixJax]
+    BackwardPassCarry = tuple[jtp.Matrix, jtp.Matrix]
     backward_pass_carry: BackwardPassCarry = (Mc, M)
 
     def backward_pass(
@@ -90,7 +90,7 @@ def crba(model: js.model.JaxSimModel, *, joint_positions: jtp.Vector) -> jtp.Mat
 
         j = i
 
-        CarryInnerFn = tuple[jtp.Int, jtp.MatrixJax, jtp.MatrixJax]
+        CarryInnerFn = tuple[jtp.Int, jtp.Matrix, jtp.Matrix]
         carry_inner_fn = (j, Fi, M)
 
         def while_loop_body(carry: CarryInnerFn) -> CarryInnerFn:

jaxsim/rbda/forward_kinematics.py

@@ -61,7 +61,7 @@ def forward_kinematics_model(
     # Propagate the kinematics
     # ========================
 
-    PropagateKinematicsCarry = tuple[jtp.MatrixJax]
+    PropagateKinematicsCarry = tuple[jtp.Matrix]
     propagate_kinematics_carry: PropagateKinematicsCarry = (W_X_i,)
 
     def propagate_kinematics(

jaxsim/rbda/jacobian.py

@@ -50,7 +50,7 @@ def jacobian(
     # Propagate kinematics
     # ====================
 
-    PropagateKinematicsCarry = tuple[jtp.MatrixJax]
+    PropagateKinematicsCarry = tuple[jtp.Matrix]
     propagate_kinematics_carry: PropagateKinematicsCarry = (i_X_0,)
 
     def propagate_kinematics(
@@ -86,9 +86,9 @@ def jacobian(
     # Checking if j ∈ κ(i) is equivalent to: κ_bool(j) is True.
     κ_bool = model.kin_dyn_parameters.support_body_array_bool[link_index]
 
-    def compute_jacobian(J: jtp.MatrixJax, i: jtp.Int) -> tuple[jtp.MatrixJax, None]:
+    def compute_jacobian(J: jtp.Matrix, i: jtp.Int) -> tuple[jtp.Matrix, None]:
 
-        def update_jacobian(J: jtp.MatrixJax, i: jtp.Int) -> jtp.MatrixJax:
+        def update_jacobian(J: jtp.Matrix, i: jtp.Int) -> jtp.Matrix:
 
             ii = i - 1
 
@@ -155,16 +155,16 @@ def jacobian_full_doubly_left(
     B_X_i = jnp.zeros(shape=(model.number_of_links(), 6, 6))
     B_X_i = B_X_i.at[0].set(jnp.eye(6))
 
-    # =============================
-    # Compute doubly-left Jacobian
-    # =============================
+    # =================================
+    # Compute doubly-left full Jacobian
+    # =================================
 
     # Allocate the Jacobian matrix.
     # The Jbb section of the doubly-left Jacobian is an identity matrix.
     J = jnp.zeros(shape=(6, 6 + model.dofs()))
     J = J.at[0:6, 0:6].set(jnp.eye(6))
 
-    ComputeFullJacobianCarry = tuple[jtp.MatrixJax, jtp.MatrixJax]
+    ComputeFullJacobianCarry = tuple[jtp.Matrix, jtp.Matrix]
     compute_full_jacobian_carry: ComputeFullJacobianCarry = (B_X_i, J)
 
     def compute_full_jacobian(
@@ -261,7 +261,7 @@ def jacobian_derivative_full_doubly_left(
     J̇ = jnp.zeros(shape=(6, 6 + model.dofs()))
 
     ComputeFullJacobianDerivativeCarry = tuple[
-        jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax
+        jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix
     ]
 
     compute_full_jacobian_derivative_carry: ComputeFullJacobianDerivativeCarry = (

jaxsim/rbda/rnea.py

@@ -132,7 +132,7 @@ def rnea(
     # Pass 1
     # ======
 
-    ForwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax, jtp.MatrixJax]
+    ForwardPassCarry = Tuple[jtp.Matrix, jtp.Matrix, jtp.Matrix, jtp.Matrix]
     forward_pass_carry: ForwardPassCarry = (v, a, i_X_0, f)
 
     def forward_pass(
@@ -186,7 +186,7 @@ def rnea(
 
     τ = jnp.zeros_like(s)
 
-    BackwardPassCarry = Tuple[jtp.MatrixJax, jtp.MatrixJax]
+    BackwardPassCarry = Tuple[jtp.Vector, jtp.Matrix]
     backward_pass_carry: BackwardPassCarry = (τ, f)
 
     def backward_pass(
@@ -201,7 +201,7 @@ def rnea(
         τ = τ.at[ii].set(τ_i.squeeze())
 
         # Propagate the force to the parent link.
-        def update_f(f: jtp.MatrixJax) -> jtp.MatrixJax:
+        def update_f(f: jtp.Matrix) -> jtp.Matrix:
 
             f_λi = f[λ[i]] + i_X_λi[i].T @ f[i]
             f = f.at[λ[i]].set(f_λi)

jaxsim/rbda/utils.py

@@ -19,7 +19,7 @@ def process_inputs(
     joint_accelerations: jtp.VectorLike | None = None,
     joint_forces: jtp.VectorLike | None = None,
     link_forces: jtp.MatrixLike | None = None,
-    standard_gravity: jtp.VectorLike | None = None,
+    standard_gravity: jtp.ScalarLike | None = None,
 ) -> tuple[
     jtp.Vector,
     jtp.Vector,

jaxsim/terrain/terrain.py

@@ -1,4 +1,7 @@
+from __future__ import annotations
+
 import abc
+import dataclasses
 
 import jax.numpy as jnp
 import jax_dataclasses
@@ -7,22 +10,23 @@ import jaxsim.typing as jtp
 
 
 class Terrain(abc.ABC):
+
     delta = 0.010
 
     @abc.abstractmethod
-    def height(self, x: float, y: float) -> float:
+    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
         pass
 
-    def normal(self, x: float, y: float) -> jtp.Vector:
+    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
         """
         Compute the normal vector of the terrain at a specific (x, y) location.
 
         Args:
-            x (float): The x-coordinate of the location.
-            y (float): The y-coordinate of the location.
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
 
         Returns:
-            jtp.Vector: The normal vector of the terrain surface at the specified location.
+            The normal vector of the terrain surface at the specified location.
         """
 
         # https://stackoverflow.com/a/5282364
@@ -40,43 +44,117 @@ class Terrain(abc.ABC):
 
 @jax_dataclasses.pytree_dataclass
 class FlatTerrain(Terrain):
-    def height(self, x: float, y: float) -> float:
-        return 0.0
+
+    z: float = dataclasses.field(default=0.0, kw_only=True)
+
+    @staticmethod
+    def build(height: jtp.FloatLike) -> FlatTerrain:
+
+        return FlatTerrain(z=float(height))
+
+    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
+
+        return jnp.array(self.z, dtype=float)
+
+    def __hash__(self) -> int:
+
+        return hash(self.z)
+
+    def __eq__(self, other: FlatTerrain) -> bool:
+
+        if not isinstance(other, FlatTerrain):
+            return False
+
+        return self.z == other.z
 
 
 @jax_dataclasses.pytree_dataclass
-class PlaneTerrain(Terrain):
-    plane_normal: list = jax_dataclasses.field(default_factory=lambda: [0, 0, 1.0])
+class PlaneTerrain(FlatTerrain):
+
+    plane_normal: tuple[float, float, float] = jax_dataclasses.field(
+        default=(0.0, 0.0, 0.0), kw_only=True
+    )
 
     @staticmethod
-    def build(plane_normal: list) -> "PlaneTerrain":
+    def build(
+        plane_normal: jtp.VectorLike, plane_height_over_origin: jtp.FloatLike = 0.0
+    ) -> PlaneTerrain:
         """
         Create a PlaneTerrain instance with a specified plane normal vector.
 
         Args:
-            plane_normal (list): The normal vector of the terrain plane.
+            plane_normal: The normal vector of the terrain plane.
+            plane_height_over_origin: The height of the plane over the origin.
 
         Returns:
             PlaneTerrain: A PlaneTerrain instance.
         """
-        if not isinstance(plane_normal, list):
-            raise TypeError(
-                f"Expected a list for the plane normal vector, got: {type(plane_normal)}."
-            )
 
-        return PlaneTerrain(plane_normal=plane_normal)
+        plane_normal = jnp.array(plane_normal, dtype=float)
+        plane_height_over_origin = jnp.array(plane_height_over_origin, dtype=float)
+
+        if plane_normal.shape != (3,):
+            msg = "Expected a 3D vector for the plane normal, got '{}'."
+            raise ValueError(msg.format(plane_normal.shape))
 
-    def height(self, x: float, y: float) -> float:
+        # Make sure that the plane normal is a unit vector.
+        plane_normal = plane_normal / jnp.linalg.norm(plane_normal)
+
+        return PlaneTerrain(
+            z=float(plane_height_over_origin),
+            plane_normal=tuple(plane_normal.tolist()),
+        )
+
+    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:
         """
         Compute the height of the terrain at a specific (x, y) location on a plane.
 
         Args:
-            x (float): The x-coordinate of the location.
-            y (float): The y-coordinate of the location.
+            x: The x-coordinate of the location.
+            y: The y-coordinate of the location.
 
         Returns:
-            float: The height of the terrain at the specified location on the plane.
+            The height of the terrain at the specified location on the plane.
         """
 
-        a, b, c = self.plane_normal
-        return -(a * x + b * y) / c
+        # Equation of the plane:      A x + B y + C z + D = 0
+        # Normal vector coordinates:  (A, B, C)
+        # The height over the origin: -D/C
+
+        # Get the plane equation coefficients from the terrain normal.
+        A, B, C = self.plane_normal
+
+        # Compute the final coefficient D considering the terrain height.
+        D = -C * self.z
+
+        # Invert the plane equation to get the height at the given (x, y) coordinates.
+        return jnp.array(-(A * x + B * y + D) / C).astype(float)
+
+    def __hash__(self) -> int:
+
+        from jaxsim.utils.wrappers import HashedNumpyArray
+
+        return hash(
+            (
+                hash(self.z),
+                HashedNumpyArray.hash_of_array(
+                    array=jnp.array(self.plane_normal, dtype=float)
+                ),
+            )
+        )
+
+    def __eq__(self, other: PlaneTerrain) -> bool:
+
+        if not isinstance(other, PlaneTerrain):
+            return False
+
+        if not (
+            jnp.allclose(self.z, other.z)
+            and jnp.allclose(
+                jnp.array(self.plane_normal, dtype=float),
+                jnp.array(other.plane_normal, dtype=float),
+            )
+        ):
+            return False
+
+        return True

jaxsim/typing.py

@@ -7,14 +7,14 @@ import jax
 # JAX types
 # =========
 
-ScalarJax = jax.Array
-IntJax = ScalarJax
-BoolJax = ScalarJax
-FloatJax = ScalarJax
+Array = jax.Array
+Scalar = Array
+Vector = Array
+Matrix = Array
 
-ArrayJax = jax.Array
-VectorJax = ArrayJax
-MatrixJax = ArrayJax
+Int = Scalar
+Bool = Scalar
+Float = Scalar
 
 PyTree = (
     dict[Hashable, "PyTree"] | list["PyTree"] | tuple["PyTree"] | None | jax.Array | Any
@@ -24,19 +24,11 @@ PyTree = (
 # Mixed JAX / NumPy types
 # =======================
 
-Array = jax.typing.ArrayLike
-Scalar = Array
-Vector = Array
-Matrix = Array
+ArrayLike = jax.typing.ArrayLike | tuple
+ScalarLike = int | float | Scalar | ArrayLike
+VectorLike = Vector | ArrayLike | tuple
+MatrixLike = Matrix | ArrayLike
 
-Int = int | IntJax
-Bool = bool | ArrayJax
-Float = float | FloatJax
-
-ScalarLike = Scalar | int | float
-ArrayLike = Array
-VectorLike = Vector
-MatrixLike = Matrix
-IntLike = Int
-BoolLike = Bool
-FloatLike = Float
+IntLike = int | Int | jax.typing.ArrayLike
+BoolLike = bool | Bool | jax.typing.ArrayLike
+FloatLike = float | Float | jax.typing.ArrayLike