加载 LeRobot 格式的 episode 数据

Load normalization statistics from checkpoint or config assets.

Load normalization statistics from checkpoint or config assets.

Main function

Test asynchronous inference server

Start server

测试 LeRobot 推理观察器

Quick smoke test entrypoint.

主函数

Map input image keys to model keys and inject a right wrist placeholder when missing.

Map input image keys to model keys (matching LiberoInputs/LiberoSubtaskInputs). The model expects: base_0_rgb, left_wrist_0_rgb, right_wrist_

Marks the end of an episode.

Normalize state using quantile stats from norm_stats and pad to expected dimension. Args: state: Raw state array (e.g., 8D for LIBERO

Normalize state using quantile stats from norm_stats and pad to expected dimension. Args: state: Raw state array (e.g., 8D for LIBERO

Called when an episode ends.

Called when an episode starts.

Append a step to the episode.

Computes sine-cosine positional embedding vectors for scalar positions.

Computes sine-cosine positional embedding vectors for scalar positions.

Follows the MoCo v3 logic.

Prepare model-ready observation data.

Temporarily prevent keyboard interrupts by delaying them until after the protected code.

Print statistics for all keys in a concise format.

Process inference request - generates subtask first, then actions. This is a simplified synchronous approach where: 1. Subtas

Process inference request - generates subtask first, then actions. This is a simplified synchronous approach where: 1. Subtas

Like np.put_along_axis(..., axis=-1), since jax is missing it.

Like np.put_along_axis(..., axis=-1), since jax is missing it.

Register new client

Save a checkpoint with model state, optimizer state, and metadata.

Serialize the running statistics to a JSON string.

Convert Gemma JAX parameters to PyTorch format.

Load and process params by restoring via JAX model loader first. This respects dtype conversions that occur during model restore.

Convert PaliGemma JAX parameters to PyTorch format.

Continuously generate action sequences

启动服务器

Start WebSocket server skip_init: If True, skip model initialization, only for connectivity testing.

测试基本功能

Test batch inference requests

Test periodic refresh functionality

Test single inference request

Transform the dataset by applying the data transforms.

Transform the dataset by applying the data transforms.

Converts a PyTree into a human-readable string for logging. Optionally, `interp_func` can be provided to convert the leaf values to more meaningfu

Fills the input Tensor with values drawn from a truncated normal distribution. The values are effectively drawn from the normal distribution :

Unnormalize model actions to action space.

Unnormalize actions using quantile stats from norm_stats. Args: actions: Normalized action array from model output norm_stats

Unregister client

Save the timings to a parquet file.

Factory function, because linen really don't like what I'm doing!

Apply transformation to the data. Args: data: The data to apply the transform to. This is a possibly nested dictionary that conta