Calculate the relative rotation matrices of N bodies with respect to a main rotation matrix. Args: main_rot (array): Rotation matrix

Calculate the relative velocity vel_a-vel_b expressed in the local frame of vel_a. Args: vel_a (array): 6D velocity vector of body A

To ensure that the data structure of the current model and the trajectory data have the same dimensionality and order for all support

Collect trajectories of given unroll_length.

Get the names of the bodies in the model. If keys is not None, only return the names of the bodies that are in keys, otherwise return

Return the gyroscope readings in the local frame.

Return the gyroscope readings in the local frame.

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

Get the names of the sites in the model. If keys is not None, only return the names of the sites that are in keys, otherwise return a

Perform interpolation for a batch of trajectories of rotation matrices. Args: xmats: Array of shape (T, N, 9) containing

Log metrics to console.

Get the distance and normal of the collision between two geoms. Args: geom_id1 (int): geom id in Mujoco model. geom_id2 (int

Get qvel index of a joint in mujoco data structure. Args: j_id (int): joint id. model (mjModel): mujoco model. Returns:

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

Plays a demo of the loaded trajectory by forcing the model positions to the ones in the trajectories at every step. Args:

Quaternion representing rotation around z-axis by yaw.

Converts a quaternion into a 9-dimensional rotation matrix.

Extract yaw (rotation around z-axis) from quaternion.

Rotate vector v by quaternion q (scalar-first convention).

Sets the Mujoco datastructure to the state specified in the trajectory data. Args: data (MjData): The Mujoco data struct

Serializes the trajectory to dict. Returns: A dictionary containing the trajectory data.

r""" Args: env: exp env traj_sample: get by `traj_handler.get_current_traj_data(traj_carry, np)`, will call Trajectory

Constructor. Args: model (mjModel): Current model. traj_path (string): path with the trajectory for the mode

Init. Args: eval_env: Batched environment to run evals on. eval_policy_fn: Function returning the policy from the policy

Helper function to perform reset and update state.

Add a new body to the trajectory model. Args: body_rootid (Array): Root id of the new body. body_weldid (Arr

Adds a new body with a default value to the trajectory data. Args: xpos_value (float): Default position value for the tr

Add a new joint to the trajectory model. Args: jnt_type: Type of the new joint. backend (Union[jax, numpy]):

Adds a new joint with a default value to the trajectory data. Args: qpos_value (float): Default position value for the t

Add a new site to the trajectory model. Args: site_pos (Array): Position of the new site. site_quat (Array):

Adds a new site with a default value for the position/velocity and an identity matrix as a rotation to the trajectory data.

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

Compute angular velocity from quaternion difference using angle-axis extraction. Quaternion convention: scalar-first (w, x, y, z).

Calculate the velocities of a batch of bodies in world frame. This function is implemented similarly to Mujoco's mj_objectVelocity function

Calculate the relative velocities of the bodies in b_ids to the main body. Args: cvel (array): Data array containing the velocities

Calculate the relative quaternions of the bodies in b_ids to the main body. Args: xquat (array): Data array containing the quaternio