jaxsim 0.1.dev401__py3-none-any.whl → 0.2.0__py3-none-any.whl

jaxsim/physics/algos/utils.py

@@ -1,69 +0,0 @@
-from typing import Tuple
-
-import jax.numpy as jnp
-
-import jaxsim.typing as jtp
-from jaxsim.physics.model.physics_model import PhysicsModel
-
-
-def process_inputs(
-    physics_model: PhysicsModel,
-    xfb: jtp.Vector | None = None,
-    q: jtp.Vector | None = None,
-    qd: jtp.Vector | None = None,
-    qdd: jtp.Vector | None = None,
-    tau: jtp.Vector | None = None,
-    f_ext: jtp.Matrix | None = None,
-) -> Tuple[jtp.Vector, jtp.Vector, jtp.Vector, jtp.Vector, jtp.Vector, jtp.Matrix]:
-    """
-    Adjust the inputs to the physics model.
-
-    Args:
-        physics_model: The physics model.
-        xfb: The variables of the base link.
-        q: The generalized coordinates.
-        qd: The generalized velocities.
-        qdd: The generalized accelerations.
-        tau: The generalized forces.
-        f_ext: The link external forces.
-
-    Returns:
-        The adjusted inputs.
-    """
-
-    # Remove extra dimensions
-    q = q.squeeze() if q is not None else jnp.zeros(physics_model.dofs())
-    qd = qd.squeeze() if qd is not None else jnp.zeros(physics_model.dofs())
-    qdd = qdd.squeeze() if qdd is not None else jnp.zeros(physics_model.dofs())
-    tau = tau.squeeze() if tau is not None else jnp.zeros(physics_model.dofs())
-    xfb = xfb.squeeze() if xfb is not None else jnp.zeros(13).at[0].set(1)
-    f_ext = (
-        f_ext.squeeze()
-        if f_ext is not None
-        else jnp.zeros(shape=(physics_model.NB, 6)).squeeze()
-    )
-
-    # Fix case with just 1 DoF
-    q = jnp.atleast_1d(q)
-    qd = jnp.atleast_1d(qd)
-    qdd = jnp.atleast_1d(qdd)
-    tau = jnp.atleast_1d(tau)
-
-    # Fix case with just 1 body
-    f_ext = jnp.atleast_2d(f_ext)
-
-    # Validate dimensions
-    dofs = physics_model.dofs()
-
-    if xfb is not None and xfb.shape[0] != 13:
-        raise ValueError(xfb.shape)
-    if q is not None and q.shape[0] != dofs:
-        raise ValueError(q.shape, dofs)
-    if qd is not None and qd.shape[0] != dofs:
-        raise ValueError(qd.shape, dofs)
-    if tau is not None and tau.shape[0] != dofs:
-        raise ValueError(tau.shape, dofs)
-    if f_ext is not None and f_ext.shape != (physics_model.NB, 6):
-        raise ValueError(f_ext.shape, (physics_model.NB, 6))
-
-    return xfb, q, qd, qdd, tau, f_ext

jaxsim/physics/model/ground_contact.py

@@ -1,55 +0,0 @@
-import dataclasses
-
-import jax.numpy as jnp
-import jax_dataclasses
-import numpy as np
-import numpy.typing as npt
-from jax_dataclasses import Static
-
-from jaxsim.parsers.descriptions import ModelDescription
-
-
-@jax_dataclasses.pytree_dataclass
-class GroundContact:
-    """
-    A class for managing collidable points in a robot model.
-
-    This class is used to store and manage information about collidable points on a robot model,
-    such as their positions and the corresponding bodies (links) they are associated with.
-
-    Attributes:
-        point (npt.NDArray): An array of shape (3, N) representing the 3D positions of collidable points.
-        body (Static[npt.NDArray]): An array of integers representing the indices of the bodies (links) associated with each collidable point.
-    """
-
-    point: npt.NDArray = dataclasses.field(default_factory=lambda: jnp.array([]))
-    body: Static[npt.NDArray] = dataclasses.field(
-        default_factory=lambda: np.array([], dtype=int)
-    )
-
-    @staticmethod
-    def build_from(
-        model_description: ModelDescription,
-    ) -> "GroundContact":
-        if len(model_description.collision_shapes) == 0:
-            return GroundContact()
-
-        # Get all the links so that we can take their updated index
-        links_dict = {link.name: link for link in model_description}
-
-        # Get all the enabled collidable points of the model
-        collidable_points = model_description.all_enabled_collidable_points()
-
-        # Build the GroundContact attributes
-        points = jnp.vstack([cp.position for cp in collidable_points]).T
-        link_index_of_points = np.array(
-            [links_dict[cp.parent_link.name].index for cp in collidable_points]
-        )
-
-        # Build the object
-        gc = GroundContact(point=points, body=link_index_of_points)
-
-        assert gc.point.shape[0] == 3
-        assert gc.point.shape[1] == len(gc.body)
-
-        return gc

jaxsim/physics/model/physics_model.py

@@ -1,358 +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
-    )
-
-    def __post_init__(self):
-        if self.initial_state is None:
-            initial_state = PhysicsModelState.zero(physics_model=self)
-            object.__setattr__(self, "initial_state", initial_state)
-
-    @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,174 +0,0 @@
-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 zero(
-        physics_model: "jaxsim.physics.model.physics_model.PhysicsModel",
-    ) -> "PhysicsModelState":
-        return PhysicsModelState(
-            joint_positions=jnp.zeros(physics_model.dofs()),
-            joint_velocities=jnp.zeros(physics_model.dofs()),
-        )
-
-    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):
-    tau: jtp.VectorJax
-    f_ext: jtp.MatrixJax
-
-    @staticmethod
-    def zero(
-        physics_model: "jaxsim.physics.model.physics_model.PhysicsModel",
-    ) -> "PhysicsModelInput":
-        ode_input = PhysicsModelInput(
-            tau=jnp.zeros(physics_model.dofs()),
-            f_ext=jnp.zeros(shape=(physics_model.NB, 6)),
-        )
-
-        assert ode_input.valid(physics_model)
-        return ode_input
-
-    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