Jax inteface

Jax inteface#

This module provides the Jax implementation of the Rigid Body Dynamics algorithms.

Tip

We suggest to jax.jit the functions as it will make them run faster!

Tip

The functions in this module can be also jax.vmap-ed to run on batches of inputs.

Note

The first time you run a jax.jit-ed function, it will take a bit longer to execute as they are being compiled by Jax.

Note

If the GPU support for JAX is needed, follow the instructions in the Jax documentation.

class adam.jax.computations.KinDynComputations(urdfstring: str, joints_name_list: list = None, root_link: str = None, gravity: array = Array([0., 0., -9.80665, 0., 0., 0.], dtype=float32))[source]#

Bases: object

This is a small class that retrieves robot quantities using Jax for Floating Base systems.

set_frame_velocity_representation(representation: Representations) → None[source]#

Sets the representation of the velocity of the frames

Parameters:

representation (Representations) – The representation of the velocity

mass_matrix(base_transform: array, joint_positions: array)[source]#

Returns the Mass Matrix functions computed the CRBA

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

Returns:

Mass Matrix

Return type:

M (jnp.array)

centroidal_momentum_matrix(base_transform: array, joint_positions: array)[source]#

Returns the Centroidal Momentum Matrix functions computed the CRBA

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

Returns:

Centroidal Momentum matrix

Return type:

Jcc (jnp.array)

relative_jacobian(frame: str, joint_positions: array)[source]#

Returns the Jacobian between the root link and a specified frame frames

Parameters:

  • frame (str) – The tip of the chain

  • joint_positions (jnp.array) – The joints position

Returns:

The Jacobian between the root and the frame

Return type:

J (jnp.array)

jacobian_dot(frame: str, base_transform: array, joint_positions: array, base_velocity: array, joint_velocities: array) → array[source]#

Returns the Jacobian derivative relative to the specified frame

Parameters:

  • frame (str) – The frame to which the jacobian will be computed

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

  • base_velocity (jnp.array) – The base velocity

  • joint_velocities (jnp.array) – The joint velocities

Returns:

The Jacobian derivative relative to the frame

Return type:

Jdot (jnp.array)

forward_kinematics(frame: str, base_transform: array, joint_positions: array)[source]#

Computes the forward kinematics relative to the specified frame

Parameters:

  • frame (str) – The frame to which the fk will be computed

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

Returns:

The fk represented as Homogenous transformation matrix

Return type:

H (jnp.array)

jacobian(frame: str, base_transform: array, joint_positions: array)[source]#

Returns the Jacobian relative to the specified frame

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • s (jnp.array) – The joints position

  • frame (str) – The frame to which the jacobian will be computed

Returns:

The Jacobian relative to the frame

Return type:

J_tot (jnp.array)

bias_force(base_transform: array, joint_positions: array, base_velocity: array, joint_velocities: array) → array[source]#

Returns the bias force of the floating-base dynamics equation, using a reduced RNEA (no acceleration and external forces)

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

  • base_velocity (jnp.array) – The base velocity

  • joint_velocities (jnp.array) – The joints velocity

Returns:

the bias force

Return type:

h (jnp.array)

coriolis_term(base_transform: array, joint_positions: array, base_velocity: array, joint_velocities: array) → array[source]#

Returns the coriolis term of the floating-base dynamics equation, using a reduced RNEA (no acceleration and external forces)

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

  • base_velocity (jnp.array) – The base velocity

  • joint_velocities (jnp.array) – The joints velocity

Returns:

the Coriolis term

Return type:

C (jnp.array)

gravity_term(base_transform: array, joint_positions: array) → array[source]#

Returns the gravity term of the floating-base dynamics equation, using a reduced RNEA (no acceleration and external forces)

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

Returns:

the gravity term

Return type:

G (jnp.array)

CoM_position(base_transform: array, joint_positions: array) → array[source]#

Returns the CoM position

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

Returns:

The CoM position

Return type:

CoM (jnp.array)

CoM_jacobian(base_transform: array, joint_positions: array) → array[source]#

Returns the CoM Jacobian

Parameters:

  • base_transform (jnp.array) – The homogenous transform from base to world frame

  • joint_positions (jnp.array) – The joints position

Returns:

The CoM Jacobian

Return type:

Jcom (jnp.array)

get_total_mass() → float[source]#

Returns the total mass of the robot

Returns:

The total mass

Return type:

mass

Contents