Convenience method for initializing all parameters, necessary due to the quirks of linen.

Update the running statistics with a batch of vectors. Args: vectors (np.ndarray): An array where all dimensions except

Create a model with the given parameters.

Send inference request Args: images: Image path dictionary, keys are image types, values are image file paths

Helper method to apply gradient checkpointing if enabled.

Close the websocket connection.

Log detailed memory usage information.

Connect to server

This method defines the mapping between unstructured data (i.e., nested dict) to the structured Observation format.

Flatten a nested dictionary. Uses '/' as the separator.

Asynchronous inference API compatible with existing call sites.

Loads the model weights. Args: params: Parameters of the model. This is a nested structure of array-like objects that

Adapted from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_

Decode tokens back to text, truncating at EOS and removing padding.

Compute and return the statistics of the vectors processed so far. Returns: dict: A dictionary containing the computed s

持续观察推理过程

Pad an array to the target dimension with zeros along the specified axis.

Append transforms to the group and return a new group. Args: inputs: Appended to the *end* of the current input transforms.

Build the full token sequence for Pi05 hierarchical training. Constructs a structured prompt that concatenates three segments in order:

Get the mtime of a given date at midnight UTC.

Used to convert between aloha and pi joint angles.

通知所有观察者

parse LeRobot image to model input format, strictly corresponding to image conversion logic

Helper method to prepare 4D attention masks for transformer.

保存观察数据

Set folder permission to be read, write and searchable.

chmod requires executable permission to be set, so we skip if the permission is already match with the target.

添加观察回调函数

Get the data config for this data loader.

Disconnect

Do a full training forward pass and compute the loss (batch_size x num_steps x num_motors)

Returns the freeze filter based on the model config.

Send an inference request and await matched response.

Asynchronously initialize model

Copied from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_a

Adapted from big_vision. Tokens can attend to valid inputs tokens which have a cumulative mask_ar smaller or equal to theirs. This way `mask_

观察单次推理

从 LeRobot 数据准备图像

Prepare observation data

从 LeRobot 数据准备状态

Record a timing measurement for the given key.

Reset the agent to its initial state.

Runs the runtime loop continuously until stop() is called or the environment is done.

Applies RoPE positions [B, L] to x [B, L, H, D].

Applies RoPE positions [B, L] to x [B, L, H, D].

Applies RoPE positions [B, L] to x [B, L, H, D].

Since we are sharing cache directory with containerized runtime as well as training script, we need to ensure that the cache directory has the cor

Applies gated residual connection with optional gate parameter. Args: x: Input tensor (residual) y: Output tensor to be

Get bins per dimension based on codebook size, from the original FSQ paper.