Calculate the difference in rigid body positions between the current data and the reference data.

Calculate the global rotation matrices of N bodies in the world frame given a main rotation matrix. Args: main_rot (array): Rotation

Calculate the relative velocity of two 6D velocities in world frame. Args: xvel (array): Data array containing the velocities of the

Concatenate a list of TrajectoryData instances given that the TrajectoryInfos are equivalent. Args: traj_datas (list): L

Collect trajectories of given unroll_length.

Return the accelerometer readings in the local frame.

Returns a list of all instances of this class.

Returns a list of all instances of this class.

Get the contact point between two geoms

Return the angular velocity of the robot in the world frame.

Return the linear velocity of the robot in the world frame.

Return the gravity vector in the world frame.

Compute linear velocity from position difference. Args: pos_curr: Current position, shape (3,) pos_last: Last position,

Loads and compiles the Mujoco model from an XML file or MjSpec object. Args: xml_file (Union[str, MjSpec]): Path to the XML file or

Check if two geoms collide. Args: geom_id1 (int): geom id in Mujoco model. geom_id2 (int): geom id in Mujoco model.

Get the joint ID (in the Mujoco datastructure) from the joint name.

Get qvel index of a joint in mujoco data structure. Args: j_name (str): joint name. model (mjModel): mujoco model. Retu

Find geom id in mujoco specification. Args: spec (MjSpec): Mujoco specification. geom_name (str): geom name. Returns:

Converts a quaternion to an angle. (quaternions is expected to be scalar last)

Return quaternion conjugate.

Flip quaternion to ensure w >= 0 (keep sign consistency). Works for shape (4,) or (B,4).

Converts a quaternion from scalar-first to scalar-last representation.

Converts a quaternion from scalar-last to scalar-first representation.

Convert a unit quaternion (scalar-first: w, x, y, z) to a 3x3 rotation matrix.

Calculate the angular distance between two rotations represented by quaternions. Args: q1 (array): First quaternion. Shape (4,) or (

Remove the bodies with the specified ids from the trajectory model. Args: body_ids (jax.Array): Array of body ids to rem

Remove the bodies with the specified ids from the trajectory data. Args: body_ids (Union[jax.Array, np.ndarray]): Array

Remove the joints with the specified ids from the trajectory model. Args: joint_ids (Union[jax.Array, np.ndarray]): Arra

Remove the joints with the specified ids from the trajectory data. Args: joint_qpos_ids (Union[jax.Array, np.ndarray]):

Remove the sites with the specified ids from the trajectory model. Args: site_ids (Union[jax.Array, np.ndarray]): Array

Remove the sites with the specified ids from the trajectory data. Args: site_ids (Union[jax.Array, np.ndarray]): Array o

Args: new_order (list[int]): List of indices of new body order. backend (ModuleType): Backend to use for the computa

Reorder the bodies in the trajectory data. Args: new_order (Union[jax.Array, np.ndarray]): Array of indices specifying t

Args: new_order (list[int]): List of indices of new joint order. backend (ModuleType): Backend to use for the comput

Reorder the joints in the trajectory data. Args: new_order_qpos (Union[jax.Array, np.ndarray]): Array of indices specify

Args: new_order (list[int]): List of indices of new site order. backend (ModuleType): Backend to use for the computa

Reorder the sites in the trajectory data. Args: new_order (Union[jax.Array, np.ndarray]): Array of indices specifying th