Set all files to be read & writable, if it is a script, keep it as a script.

Invalidate the cache if it is expired. Return True if the cache was invalidated.

A single step of the runtime loop.

Update KV cache with new key-value pairs (single token). Args: k: New key tensor of shape (batch, num_heads, 1, head_dim)

Update KV cache with new values

Update KV cache with new values

Update histograms with new vectors.

Take an action in the environment.

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

Batch inference requests

Close the websocket connection.

Convert PI0 JAX checkpoint to PyTorch format. Args: checkpoint_dir: Path to the JAX checkpoint precision: Model precision (f

Create a default policy for the given environment.

创建模拟的 LeRobot 数据

Create a policy from the given arguments.

Create an RLDS data loader for training. Note: This data loader requires some extra dependencies -- see examples/droid/README_train.md Args:

Computes sine-cosine positional embedding vectors for scalar positions.

创建测试用的图像文件

Create a data loader for training. Args: data_config: The data configuration. action_horizon: The action horizon. batch_s

Create a dataset for training.

Converts a string like "B" or "B/32" into a params dict.

Apply one denoising step of the noise `x_t` at a given timestep.

Deserialize the running statistics from a JSON string.

Drain queued periodic subtask refresh messages.

使用示例

Generate action trajectory from observation.

Generate low-level subtask text from observation.

Query the agent for the next action.

Returns config for specified gemma variant.

r""" Returns: image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by

Get the latest checkpoint step number from a checkpoint directory.

Query the environment for the current state.

Get a safe dtype for the given device type.

Get statistics for the given key.

Enable gradient checkpointing for memory optimization.

Infer actions from observations.

Asynchronous inference function, supports subtask generation and periodic refresh Args: images: Image dictionary, keys a

Send inference request and receive subtask + actions. Args: images: Dictionary of image arrays high_level_pro

Send inference request and receive subtask + actions. Args: images: Dictionary of image arrays high_level_pro

Custom logging format for better readability.

Initialize wandb logging.

This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This m

Allow the environment to signal that the episode is complete. This will be called after each step. It should return `True` if the episode is

Converts input from left-align to right-aligned.

Converts input from left-align to right-aligned.

Listen for periodic refresh messages

Load the latest checkpoint and return the global step.

Load image

Load image from path with random fallback.

Load image, supports fallback to random image

Load image, fallback to random image when loading fails.

Load JAX model from checkpoint and print all parameter keys. Args: checkpoint_dir: Path to the checkpoint directory