jaxsim 0.2.dev191__py3-none-any.whl → 0.2.dev364__py3-none-any.whl

jaxsim/physics/model/physics_model.py

@@ -1,388 +0,0 @@
-import dataclasses
-from typing import Dict, Union
-
-import jax.lax
-import jax.numpy as jnp
-import jax_dataclasses
-import numpy as np
-from jax_dataclasses import Static
-
-import jaxsim.parsers
-import jaxsim.physics
-import jaxsim.typing as jtp
-from jaxsim.parsers.descriptions import JointDescriptor, JointType
-from jaxsim.physics import default_gravity
-from jaxsim.sixd import se3
-from jaxsim.utils import JaxsimDataclass, not_tracing
-
-from .ground_contact import GroundContact
-from .physics_model_state import PhysicsModelState
-
-
-@jax_dataclasses.pytree_dataclass
-class PhysicsModel(JaxsimDataclass):
-    """
-    A read-only class to store all the information necessary to run RBDAs on a model.
-
-    This class contains information about the physics model, including the number of bodies, initial state, gravity,
-    floating base configuration, ground contact points, and more.
-
-    Attributes:
-        NB (Static[int]): The number of bodies in the physics model.
-        initial_state (PhysicsModelState): The initial state of the physics model (default: None).
-        gravity (jtp.Vector): The gravity vector (default: [0, 0, 0, 0, 0, 0]).
-        is_floating_base (Static[bool]): A flag indicating whether the model has a floating base (default: False).
-        gc (GroundContact): The ground contact points of the model (default: empty GroundContact instance).
-        description (Static[jaxsim.parsers.descriptions.model.ModelDescription]): A description of the model (default: None).
-    """
-
-    NB: Static[int]
-    initial_state: PhysicsModelState = dataclasses.field(default=None)
-    gravity: jtp.Vector = dataclasses.field(
-        default_factory=lambda: jnp.hstack(
-            [np.zeros(3), jaxsim.physics.default_gravity()]
-        )
-    )
-    is_floating_base: Static[bool] = dataclasses.field(default=False)
-    gc: GroundContact = dataclasses.field(default_factory=lambda: GroundContact())
-    description: Static[jaxsim.parsers.descriptions.model.ModelDescription] = (
-        dataclasses.field(default=None)
-    )
-
-    _parent_array_dict: Static[Dict[int, int]] = dataclasses.field(default_factory=dict)
-    _jtype_dict: Static[Dict[int, Union[JointType, JointDescriptor]]] = (
-        dataclasses.field(default_factory=dict)
-    )
-    _tree_transforms_dict: Dict[int, jtp.Matrix] = dataclasses.field(
-        default_factory=dict
-    )
-    _link_inertias_dict: Dict[int, jtp.Matrix] = dataclasses.field(default_factory=dict)
-
-    _joint_friction_static: Dict[int, float] = dataclasses.field(default_factory=dict)
-    _joint_friction_viscous: Dict[int, float] = dataclasses.field(default_factory=dict)
-
-    _joint_limit_spring: Dict[int, float] = dataclasses.field(default_factory=dict)
-    _joint_limit_damper: Dict[int, float] = dataclasses.field(default_factory=dict)
-
-    _joint_motor_inertia: Dict[int, float] = dataclasses.field(default_factory=dict)
-    _joint_motor_gear_ratio: Dict[int, float] = dataclasses.field(default_factory=dict)
-    _joint_motor_viscous_friction: Dict[int, float] = dataclasses.field(
-        default_factory=dict
-    )
-
-    _link_masses: jtp.Vector = dataclasses.field(init=False)
-    _link_spatial_inertias: jtp.Vector = dataclasses.field(init=False)
-    _joint_position_limits_min: jtp.Matrix = dataclasses.field(init=False)
-    _joint_position_limits_max: jtp.Matrix = dataclasses.field(init=False)
-
-    def __post_init__(self):
-        if self.initial_state is None:
-            initial_state = PhysicsModelState.zero(physics_model=self)
-            object.__setattr__(self, "initial_state", initial_state)
-
-        ordered_links = sorted(
-            list(self.description.links_dict.values()),
-            key=lambda l: l.index,
-        )
-
-        ordered_joints = sorted(
-            list(self.description.joints_dict.values()),
-            key=lambda j: j.index,
-        )
-
-        from jaxsim.utils import Mutability
-
-        with self.mutable_context(
-            mutability=Mutability.MUTABLE_NO_VALIDATION, restore_after_exception=False
-        ):
-            self._link_masses = jnp.stack([link.mass for link in ordered_links])
-            self._link_spatial_inertias = jnp.stack(
-                [self._link_inertias_dict[l.index] for l in ordered_links]
-            )
-
-            s_min = jnp.array([j.position_limit[0] for j in ordered_joints])
-            s_max = jnp.array([j.position_limit[1] for j in ordered_joints])
-            self._joint_position_limits_min = jnp.vstack([s_min, s_max]).min(axis=0)
-            self._joint_position_limits_max = jnp.vstack([s_min, s_max]).max(axis=0)
-
-    @staticmethod
-    def build_from(
-        model_description: jaxsim.parsers.descriptions.model.ModelDescription,
-        gravity: jtp.Vector = default_gravity(),
-    ) -> "PhysicsModel":
-        if gravity.size != 3:
-            raise ValueError(gravity.size)
-
-        # Currently, we assume that the link frame matches the frame of its parent joint
-        for l in model_description:
-            if not jnp.allclose(l.pose, jnp.eye(4)):
-                raise ValueError(f"Link '{l.name}' has unsupported pose:\n{l.pose}")
-
-        # ===================================
-        # Initialize physics model parameters
-        # ===================================
-
-        # Get the number of bodies, including the base link
-        num_of_bodies = len(model_description)
-
-        # Build the parent array λ of the floating-base model.
-        # Note: the parent of the base link is not set since it's not defined.
-        parent_array_dict = {
-            link.index: link.parent.index
-            for link in model_description
-            if link.parent is not None
-        }
-
-        # Get the 6D inertias of all links
-        link_spatial_inertias_dict = {
-            link.index: link.inertia for link in iter(model_description)
-        }
-
-        # Dict from the joint index to its type.
-        # Note: the joint index is equal to its child link index.
-        joint_types_dict = {
-            joint.index: joint.jtype for joint in model_description.joints
-        }
-
-        # Dicts from the joint index to the static and viscous friction.
-        # Note: the joint index is equal to its child link index.
-        joint_friction_static = {
-            joint.index: jnp.array(joint.friction_static, dtype=float)
-            for joint in model_description.joints
-        }
-        joint_friction_viscous = {
-            joint.index: jnp.array(joint.friction_viscous, dtype=float)
-            for joint in model_description.joints
-        }
-
-        # Dicts from the joint index to the spring and damper joint limits parameters.
-        # Note: the joint index is equal to its child link index.
-        joint_limit_spring = {
-            joint.index: jnp.array(joint.position_limit_spring, dtype=float)
-            for joint in model_description.joints
-        }
-        joint_limit_damper = {
-            joint.index: jnp.array(joint.position_limit_damper, dtype=float)
-            for joint in model_description.joints
-        }
-
-        # Dicts from the joint index to the motor inertia, gear ratio and viscous friction.
-        # Note: the joint index is equal to its child link index.
-        joint_motor_inertia = {
-            joint.index: jnp.array(joint.motor_inertia, dtype=float)
-            for joint in model_description.joints
-        }
-        joint_motor_gear_ratio = {
-            joint.index: jnp.array(joint.motor_gear_ratio, dtype=float)
-            for joint in model_description.joints
-        }
-        joint_motor_viscous_friction = {
-            joint.index: jnp.array(joint.motor_viscous_friction, dtype=float)
-            for joint in model_description.joints
-        }
-
-        # Transform between model's root and model's base link
-        # (this is just the pose of the base link in the SDF description)
-        base_link = model_description.links_dict[model_description.link_names()[0]]
-        R_H_B = model_description.transform(name=base_link.name)
-        tree_transform_0 = se3.SE3.from_matrix(matrix=R_H_B).adjoint()
-
-        # Helper to compute the transform pre(i)_H_λ(i).
-        # Given a joint 'i', it is the coordinate transform between its predecessor
-        # frame [pre(i)] and the frame of its parent link [λ(i)].
-        prei_H_λi = lambda j: model_description.relative_transform(
-            relative_to=j.name, name=j.parent.name
-        )
-
-        # Compute the tree transforms: pre(i)_X_λ(i).
-        # Given a joint 'i', it is the coordinate transform between its predecessor
-        # frame [pre(i)] and the frame of its parent link [λ(i)].
-        tree_transforms_dict = {
-            0: tree_transform_0,
-            **{
-                j.index: se3.SE3.from_matrix(matrix=prei_H_λi(j)).adjoint()
-                for j in model_description.joints
-            },
-        }
-
-        # =======================
-        # Build the initial state
-        # =======================
-
-        # Initial joint positions
-        q0 = jnp.array(
-            [
-                model_description.joints_dict[j.name].initial_position
-                for j in model_description.joints
-            ]
-        )
-
-        # Build the initial state
-        initial_state = PhysicsModelState(
-            joint_positions=q0,
-            joint_velocities=jnp.zeros_like(q0),
-            base_position=model_description.root_pose.root_position,
-            base_quaternion=model_description.root_pose.root_quaternion,
-        )
-
-        # =======================
-        # Build the physics model
-        # =======================
-
-        # Initialize the model
-        physics_model = PhysicsModel(
-            NB=num_of_bodies,
-            initial_state=initial_state,
-            _parent_array_dict=parent_array_dict,
-            _jtype_dict=joint_types_dict,
-            _tree_transforms_dict=tree_transforms_dict,
-            _link_inertias_dict=link_spatial_inertias_dict,
-            _joint_friction_static=joint_friction_static,
-            _joint_friction_viscous=joint_friction_viscous,
-            _joint_limit_spring=joint_limit_spring,
-            _joint_limit_damper=joint_limit_damper,
-            _joint_motor_gear_ratio=joint_motor_gear_ratio,
-            _joint_motor_inertia=joint_motor_inertia,
-            _joint_motor_viscous_friction=joint_motor_viscous_friction,
-            gravity=jnp.hstack([gravity.squeeze(), np.zeros(3)]),
-            is_floating_base=True,
-            gc=GroundContact.build_from(model_description=model_description),
-            description=model_description,
-        )
-
-        # Floating-base models
-        if not model_description.fixed_base:
-            return physics_model
-
-        # Fixed-base models
-        with jax_dataclasses.copy_and_mutate(physics_model) as physics_model_fixed:
-            physics_model_fixed.is_floating_base = False
-
-        return physics_model_fixed
-
-    def dofs(self) -> int:
-        return len(list(self._jtype_dict.keys()))
-
-    def set_gravity(self, gravity: jtp.Vector) -> None:
-        gravity = gravity.squeeze()
-
-        if gravity.size == 3:
-            self.gravity = jnp.hstack([gravity, 0, 0, 0])
-
-        elif gravity.size == 6:
-            self.gravity = gravity
-
-        else:
-            raise ValueError(gravity.shape)
-
-    @property
-    def parent(self) -> jtp.Vector:
-        return self.parent_array()
-
-    def parent_array(self) -> jtp.Vector:
-        """Returns λ(i)"""
-        return jnp.array([-1] + list(self._parent_array_dict.values()), dtype=int)
-
-    def support_body_array(self, body_index: jtp.Int) -> jtp.Vector:
-        """Returns κ(i)"""
-
-        κ_bool = self.support_body_array_bool(body_index=body_index)
-        return jnp.array(jnp.where(κ_bool)[0], dtype=int)
-
-    def support_body_array_bool(self, body_index: jtp.Int) -> jtp.Vector:
-        active_link = body_index
-        κ_bool = jnp.zeros(self.NB, dtype=bool)
-
-        for i in np.flip(np.arange(start=0, stop=self.NB)):
-            κ_bool, active_link = jax.lax.cond(
-                pred=(i == active_link),
-                false_fun=lambda: (κ_bool, active_link),
-                true_fun=lambda: (
-                    κ_bool.at[active_link].set(True),
-                    self.parent[active_link],
-                ),
-            )
-
-        return κ_bool
-
-    @property
-    def tree_transforms(self) -> jtp.Array:
-        X_tree = jnp.array(
-            [
-                self._tree_transforms_dict.get(idx, jnp.eye(6))
-                for idx in np.arange(start=0, stop=self.NB)
-            ]
-        )
-
-        return X_tree
-
-    @property
-    def spatial_inertias(self) -> jtp.Array:
-        M_links = jnp.array(
-            [
-                self._link_inertias_dict.get(idx, jnp.zeros(6))
-                for idx in np.arange(start=0, stop=self.NB)
-            ]
-        )
-
-        return M_links
-
-    def jtype(self, joint_index: int) -> JointType:
-        if joint_index == 0 or joint_index >= self.NB:
-            raise ValueError(joint_index)
-
-        return self._jtype_dict[joint_index]
-
-    def joint_transforms(self, q: jtp.Vector) -> jtp.Array:
-        from jaxsim.math.joint import jcalc
-
-        if not_tracing(q):
-            if q.shape[0] != self.dofs():
-                raise ValueError(q.shape)
-
-        Xj = jnp.stack(
-            [jnp.zeros(shape=(6, 6))]
-            + [
-                jcalc(jtyp=self.jtype(index + 1), q=joint_position)[0]
-                for index, joint_position in enumerate(q)
-            ]
-        )
-
-        return Xj
-
-    def motion_subspaces(self, q: jtp.Vector) -> jtp.Array:
-        from jaxsim.math.joint import jcalc
-
-        if not_tracing(var=q):
-            if q.shape[0] != self.dofs():
-                raise ValueError(q.shape)
-
-        SS = jnp.stack(
-            [jnp.vstack(jnp.zeros(6))]
-            + [
-                jcalc(jtyp=self.jtype(index + 1), q=joint_position)[1]
-                for index, joint_position in enumerate(q)
-            ]
-        )
-
-        return SS
-
-    def __eq__(self, other: "PhysicsModel") -> bool:
-        same = True
-        same = same and self.NB == other.NB
-        same = same and np.allclose(self.gravity, other.gravity)
-
-        return same
-
-    def __hash__(self):
-        return hash(self.__repr__())
-
-    def __repr__(self) -> str:
-        attributes = [
-            f"dofs: {self.dofs()},",
-            f"links: {self.NB},",
-            f"floating_base: {self.is_floating_base},",
-        ]
-        attributes_string = "\n    ".join(attributes)
-
-        return f"{type(self).__name__}(\n    {attributes_string}\n)"

jaxsim/physics/model/physics_model_state.py

@@ -1,283 +0,0 @@
-from typing import Union
-
-import jax.numpy as jnp
-import jax_dataclasses
-
-import jaxsim.physics.model.physics_model
-import jaxsim.typing as jtp
-from jaxsim.utils import JaxsimDataclass
-
-
-@jax_dataclasses.pytree_dataclass
-class PhysicsModelState(JaxsimDataclass):
-    """
-    A class representing the state of a physics model.
-
-    This class stores the joint positions, joint velocities, and the base state (position, orientation, linear velocity,
-    and angular velocity) of a physics model.
-
-    Attributes:
-        joint_positions (jtp.Vector): An array representing the joint positions.
-        joint_velocities (jtp.Vector): An array representing the joint velocities.
-        base_position (jtp.Vector): An array representing the base position (default: zeros).
-        base_quaternion (jtp.Vector): An array representing the base quaternion (default: [1.0, 0, 0, 0]).
-        base_linear_velocity (jtp.Vector): An array representing the base linear velocity (default: zeros).
-        base_angular_velocity (jtp.Vector): An array representing the base angular velocity (default: zeros).
-    """
-
-    # Joint state
-    joint_positions: jtp.Vector
-    joint_velocities: jtp.Vector
-
-    # Base state
-    base_position: jtp.Vector = jax_dataclasses.field(
-        default_factory=lambda: jnp.zeros(3)
-    )
-    base_quaternion: jtp.Vector = jax_dataclasses.field(
-        default_factory=lambda: jnp.array([1.0, 0, 0, 0])
-    )
-    base_linear_velocity: jtp.Vector = jax_dataclasses.field(
-        default_factory=lambda: jnp.zeros(3)
-    )
-    base_angular_velocity: jtp.Vector = jax_dataclasses.field(
-        default_factory=lambda: jnp.zeros(3)
-    )
-
-    @staticmethod
-    def build(
-        joint_positions: jtp.Vector | None = None,
-        joint_velocities: jtp.Vector | None = None,
-        base_position: jtp.Vector | None = None,
-        base_quaternion: jtp.Vector | None = None,
-        base_linear_velocity: jtp.Vector | None = None,
-        base_angular_velocity: jtp.Vector | None = None,
-        number_of_dofs: jtp.Int | None = None,
-    ) -> "PhysicsModelState":
-        """"""
-
-        joint_positions = (
-            joint_positions
-            if joint_positions is not None
-            else jnp.zeros(number_of_dofs)
-        )
-
-        joint_velocities = (
-            joint_velocities
-            if joint_velocities is not None
-            else jnp.zeros(number_of_dofs)
-        )
-
-        base_position = base_position if base_position is not None else jnp.zeros(3)
-
-        base_quaternion = (
-            base_quaternion
-            if base_quaternion is not None
-            else jnp.array([1.0, 0, 0, 0])
-        )
-
-        base_linear_velocity = (
-            base_linear_velocity if base_linear_velocity is not None else jnp.zeros(3)
-        )
-
-        base_angular_velocity = (
-            base_angular_velocity if base_angular_velocity is not None else jnp.zeros(3)
-        )
-
-        physics_model_state = PhysicsModelState(
-            joint_positions=jnp.array(joint_positions, dtype=float),
-            joint_velocities=jnp.array(joint_velocities, dtype=float),
-            base_position=jnp.array(base_position, dtype=float),
-            base_quaternion=jnp.array(base_quaternion, dtype=float),
-            base_linear_velocity=jnp.array(base_linear_velocity, dtype=float),
-            base_angular_velocity=jnp.array(base_angular_velocity, dtype=float),
-        )
-
-        return physics_model_state
-
-    @staticmethod
-    def build_from_physics_model(
-        joint_positions: jtp.Vector | None = None,
-        joint_velocities: jtp.Vector | None = None,
-        base_position: jtp.Vector | None = None,
-        base_quaternion: jtp.Vector | None = None,
-        base_linear_velocity: jtp.Vector | None = None,
-        base_angular_velocity: jtp.Vector | None = None,
-        physics_model: Union[
-            "jaxsim.physics.model.physics_model.PhysicsModel", None
-        ] = None,
-    ) -> "PhysicsModelState":
-        """"""
-
-        return PhysicsModelState.build(
-            joint_positions=joint_positions,
-            joint_velocities=joint_velocities,
-            base_position=base_position,
-            base_quaternion=base_quaternion,
-            base_linear_velocity=base_linear_velocity,
-            base_angular_velocity=base_angular_velocity,
-            number_of_dofs=physics_model.dofs(),
-        )
-
-    @staticmethod
-    def zero(
-        physics_model: "jaxsim.physics.model.physics_model.PhysicsModel",
-    ) -> "PhysicsModelState":
-        return PhysicsModelState.build_from_physics_model(physics_model=physics_model)
-
-    def position(self) -> jtp.Vector:
-        return jnp.hstack(
-            [self.base_position, self.base_quaternion, self.joint_positions]
-        )
-
-    def velocity(self) -> jtp.Vector:
-        # W_v_WB: inertial-fixed representation of the base velocity
-        return jnp.hstack(
-            [
-                self.base_linear_velocity,
-                self.base_angular_velocity,
-                self.joint_velocities,
-            ]
-        )
-
-    def xfb(self) -> jtp.Vector:
-        return jnp.hstack(
-            [
-                self.base_quaternion,
-                self.base_position,
-                self.base_angular_velocity,
-                self.base_linear_velocity,
-            ]
-        )
-
-    def valid(
-        self, physics_model: "jaxsim.physics.model.physics_model.PhysicsModel"
-    ) -> bool:
-        from jaxsim.simulation.utils import check_valid_shape
-
-        valid = True
-
-        valid = check_valid_shape(
-            what="joint_positions",
-            shape=self.joint_positions.shape,
-            expected_shape=(physics_model.dofs(),),
-            valid=valid,
-        )
-
-        valid = check_valid_shape(
-            what="joint_velocities",
-            shape=self.joint_velocities.shape,
-            expected_shape=(physics_model.dofs(),),
-            valid=valid,
-        )
-
-        valid = check_valid_shape(
-            what="base_position",
-            shape=self.base_position.shape,
-            expected_shape=(3,),
-            valid=valid,
-        )
-
-        valid = check_valid_shape(
-            what="base_quaternion",
-            shape=self.base_quaternion.shape,
-            expected_shape=(4,),
-            valid=valid,
-        )
-
-        valid = check_valid_shape(
-            what="base_linear_velocity",
-            shape=self.base_linear_velocity.shape,
-            expected_shape=(3,),
-            valid=valid,
-        )
-
-        valid = check_valid_shape(
-            what="base_angular_velocity",
-            shape=self.base_angular_velocity.shape,
-            expected_shape=(3,),
-            valid=valid,
-        )
-
-        return valid
-
-
-@jax_dataclasses.pytree_dataclass
-class PhysicsModelInput(JaxsimDataclass):
-    """
-    A class representing the input to a physics model.
-
-    This class stores the joint torques and external forces acting on the bodies of a physics model.
-
-    Attributes:
-        tau: An array representing the joint torques.
-        f_ext: A matrix representing the external forces acting on the bodies of the physics model.
-    """
-
-    tau: jtp.VectorJax
-    f_ext: jtp.MatrixJax
-
-    @staticmethod
-    def build(
-        tau: jtp.VectorJax | None = None,
-        f_ext: jtp.MatrixJax | None = None,
-        number_of_dofs: jtp.Int | None = None,
-        number_of_links: jtp.Int | None = None,
-    ) -> "PhysicsModelInput":
-        """"""
-
-        tau = tau if tau is not None else jnp.zeros(number_of_dofs)
-        f_ext = f_ext if f_ext is not None else jnp.zeros(shape=(number_of_links, 6))
-
-        return PhysicsModelInput(
-            tau=jnp.array(tau, dtype=float), f_ext=jnp.array(f_ext, dtype=float)
-        )
-
-    @staticmethod
-    def build_from_physics_model(
-        tau: jtp.VectorJax | None = None,
-        f_ext: jtp.MatrixJax | None = None,
-        physics_model: Union[
-            "jaxsim.physics.model.physics_model.PhysicsModel", None
-        ] = None,
-    ) -> "PhysicsModelInput":
-        return PhysicsModelInput.build(
-            tau=tau,
-            f_ext=f_ext,
-            number_of_dofs=physics_model.dofs(),
-            number_of_links=physics_model.NB,
-        )
-
-    @staticmethod
-    def zero(
-        physics_model: "jaxsim.physics.model.physics_model.PhysicsModel",
-    ) -> "PhysicsModelInput":
-        return PhysicsModelInput.build_from_physics_model(physics_model=physics_model)
-
-    def replace(self, validate: bool = True, **kwargs) -> "PhysicsModelInput":
-        with jax_dataclasses.copy_and_mutate(self, validate=validate) as updated_input:
-            _ = [updated_input.__setattr__(k, v) for k, v in kwargs.items()]
-
-        return updated_input
-
-    def valid(
-        self, physics_model: "jaxsim.physics.model.physics_model.PhysicsModel"
-    ) -> bool:
-        from jaxsim.simulation.utils import check_valid_shape
-
-        valid = True
-
-        valid = check_valid_shape(
-            what="tau",
-            shape=self.tau.shape,
-            expected_shape=(physics_model.dofs(),),
-            valid=valid,
-        )
-
-        valid = check_valid_shape(
-            what="f_ext",
-            shape=self.f_ext.shape,
-            expected_shape=(physics_model.NB, 6),
-            valid=valid,
-        )
-
-        return valid

jaxsim/simulation/__init__.py

@@ -1,4 +0,0 @@
-from . import integrators, ode, ode_data, simulator
-from .ode_data import ODEInput, ODEState
-from .ode_integration import IntegratorType
-from .simulator import JaxSim, SimulatorData