decode token ids to continuous actions. action_token_id: np.ndarray, (n, 7), token ids. return: np.ndarray, (n, 7), continuous actions

Converts relative gripper actions (+1 for closing, -1 for opening) to absolute actions (0 = closed; 1 = open). Assumes that the first relati

This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, he

Converts gripper actions from continuous to binary values (0 and 1). We exploit that fact that most of the time, the gripper is fully open (

Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_l

Allocates an integer number of threads across datasets based on weights. The final array sums to `n`, but each element is no less than 1. If

This function is responsible for loading a specific RLDS dataset from storage and getting it into a standardized format. Yields a dataset of

Randomly swaps the two exterior images (for training with single exterior input).

Relabels actions to use reached proprioceptive state; discards last timestep (no-action).

Rotates half the hidden dims of the input.

Initialize slurm distributed training environment. If argument ``port`` is not specified, then the master port will be system environment var

init with small weights to maintain stable training.

Applies common transforms that happen at a frame level. These transforms are usually more CPU-intensive, (e.g., decoding or resizing images).

Optionally applied *per-dataset* transforms that happen at a frame level. Args: chunk_filter_fn (callable, optional): Filter functio

Applies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): T

Applies common transforms that happen at a trajectory level. Such transforms are usually some sort of "relabeling" (e.g., filtering, chunking

Backproject depth map to 3D points in camera coordinate. Args: K: camera intrinsic matrix (b 3 3) depth: dept

This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please refer to the docstring o

Creates an interleaved dataset from list of dataset configs (kwargs). Returns a dataset of batched frames. Args: dataset_kwargs_list

This method forwards all its arguments to GemmaTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to the docstring of this

Embeds x to (x, sin(2^k x), cos(2^k x), ...) Different from the paper, "x" is also in the output See https://github.com/bmild

Either computes the statistics of a dataset or loads them from a cache file if this function has been called before with the same `hash_depen

Generates dataset kwargs for a given dataset mix from the Open X-Embodiment dataset. The returned kwargs (per-dataset configs) and weights ca

Generates config (kwargs) for given dataset from Open-X Embodiment.

https://github.com/huggingface/transformers/blob/v4.45.2/src/transformers/models/zoedepth/image_processing_zoedepth.py

Converts rotation matrix to R6 rotation representation (first two rows in rotation matrix). Args: mat: rotation matrix Returns:

Saves a `ds_stats.json` file.

gs_params0, gs_params1: Dict embeddings: tensor (S,E)

Split a list of indices into `chunks` chunks of roughly equal lengths.

Translates velocity actions (translation + rotation) from base frame of the robot to wrist frame. Args: velocity: 6d velocity action

Discretize continuous actions to tokens. action: np.ndarray, (n, 7), continuous actions in Cartesian or Spherical coordinates. return:

Discretize continuous actions to tokens. action: np.ndarray, (n, 3), continuous actions in Cartesian or Spherical coordinates. return:

Discretize continuous actions to tokens. action: np.ndarray, (n, 3), continuous actions in Cartesian or Spherical coordinates. return:

Discretize continuous actions to tokens. action: np.ndarray, (n,), continuous actions in Cartesian or Spherical coordinates. return: n

Discretize continuous actions to tokens. action: np.ndarray, (n, 7), continuous actions in Cartesian coordinates. return: np.ndarray,

set bin_policy if exist, otherwise, caculate bin_policy from gs_params or use uniform bin grids. gs_params: Optional[Dict], bin_policy