Retrieve the corresponding data for a given trajectory index and sub-trajectory index. Args: traj_index: Index of the tr

Helper function to extract a single attribute slice.

Helper function to extract a single attribute.

Gets sensor data given sensor name.

Loads a trajectory from a npz file. Args: path (str): Path to the trajectory to load. backend: Backend to us

Inverse of a quaternion (scalar-first convention: w, x, y, z). Assumes unit quaternion.

Multiply two quaternions (scalar-first convention: w, x, y, z).

Ensures the tensor has at least 3 dimensions by adding axes at the front if necessary. Args: tensor: Input tensor (numpy or JAX arra

Recalculate the angular velocity of the trajectory. qvel: free joint 3 linear (global) + 3 angular (local), hinge joint 1xn

Recalculate the joint velocity of the trajectory. qvel: free joint 3 linear (global) + 3 angular (local), hinge joint 1xn

Recalculate the linear velocity of the trajectory. qvel: free joint 3 linear (global) + 3 angular (local), hinge joint 1xn

Serializes the trajectory and saves it to a npz file. Args: path (str): Path to save the trajectory.

Perform SLERP interpolation for a batch of quaternions. Args: quats: Array of shape (T, 4) containing quaternions. (quat

Add a new body to the trajectory info. Args: body_name (list(str)): Body name to add. body_rootid (int): Roo

Add a new joint to the trajectory info. Args: joint_name (list(str)): Joint name to add. joint_type (mujoco.

Add a new site to the trajectory info. Args: site_name (list[str]): site name to add. site_pos (Array): Posi

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

Interpolate between two qpos configurations, handling yaw wrapping. Args: q_start (np.ndarray): Starting qpos (n,). q_end (n

Interpolate the trajectories to a new frequency. Args: traj_data: TrajectoryData instance containing the trajectories to interpolate

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

Modifies the Mujoco specification options. Args: spec (MjSpec): The Mujoco specification. option_config (Dict): Dictionary o

Remove the bodies with the specified ids from the trajectory info. Args: body_names (list[str]): List of body ids to rem

Remove the joints with the specified ids from the trajectory info. Args: joint_names (list[str]): List of joint names to

Remove the sites with the specified ids from the trajectory info. Args: site_names (list[str]): List of site ids to remo

Args: new_order (list[int]): List of indices of new body order. backend (Union[jax, numpy]): Backend to use for the

Args: new_order (list[int]): List of indices of new joint order. backend (Union[jax, numpy]): Backend to use for the

Args: new_order (list[int]): List of indices of new site order. backend (Union[jax, numpy]): Backend to use for the

Calculate the rotation error between two rotation matrices (from target to current!!!) under small angle approximation. Mathematical Princip

r""" Spherical linear interpolation between two rotation matrices.

Transforms a motion vector from one frame to another. This function is implemented similarly to Mujoco's mju_transformSpatial function, but i

r""" This function is only called in play. `next_subtraj_step_no` need to be set to `fixed_start_conf[1]` since the reset logic in play

r""" Aligns the stance foot between start and end configurations, then generates an IK trajectory for the swing foot to transition from start

Wraps the environment for training/eval if wrap_env is True.

Reset the random target state for excluded joints.

Update random targets for excluded joints if needed.

Collect data.

Collect data.

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