opentau 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

opentau/utils/transformers_patch.py

@@ -12,37 +12,268 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Module for patching transformers metadata lookup.
+"""Module for patching transformers
 
-This module monkey patches `importlib.metadata.distribution` to redirect
-'transformers' lookups to 'opentau-transformers'. This ensures that the custom
-transformers fork is correctly recognized by libraries checking for transformers
-installation.
+Most patches come from the branch fix/lerobot-openpi
 """
 
-import importlib.metadata
+from typing import Optional, Tuple
 
-# Keep a reference to the original distribution function
-_orig_distribution = importlib.metadata.distribution
+import torch
+from torch import nn
+from transformers.models.gemma import modeling_gemma
+from transformers.models.gemma.configuration_gemma import GemmaConfig
+from transformers.models.paligemma.modeling_paligemma import PaliGemmaModel
 
+# Monkey patch __init__ of GemmaConfig to fix or modify its behavior as needed.
 
-def distribution(distribution_name):
-    """Monkey patch to redirect 'transformers' metadata lookups to 'opentau-transformers'.
+_original_gemma_config_init = GemmaConfig.__init__
 
-    This function intercepts calls to `importlib.metadata.distribution`. If the
-    requested distribution is "transformers", it redirects the lookup to
-    "opentau-transformers". Otherwise, it delegates to the original implementation.
+
+def patched_gemma_config_init(
+    self, *args, use_adarms: bool = False, adarms_cond_dim: Optional[int] = None, **kwargs
+):
+    """Initializes the GemmaConfig with added ADARMS support.
+
+    Args:
+        self: The GemmaConfig instance.
+        *args: Variable length argument list.
+        use_adarms: Whether to use Adaptive RMS normalization.
+        adarms_cond_dim: The dimension of the conditioning vector for ADARMS.
+        **kwargs: Arbitrary keyword arguments.
+    """
+    # Call the original init with all other arguments
+    _original_gemma_config_init(self, *args, **kwargs)
+
+    # Initialize custom attributes
+    self.use_adarms = use_adarms
+    self.adarms_cond_dim = adarms_cond_dim
+
+    # Set default for adarms_cond_dim if use_adarms is True
+    if self.use_adarms and self.adarms_cond_dim is None:
+        # hidden_size is set by _original_gemma_config_init
+        self.adarms_cond_dim = self.hidden_size
+
+
+GemmaConfig.__init__ = patched_gemma_config_init
+
+
+# --- Modeling Patches ---
+
+
+def _gated_residual(x, y, gate):
+    """
+    Applies gated residual connection with optional gate parameter.
+
+    Args:
+        x: Input tensor (residual)
+        y: Output tensor to be added
+        gate: Optional gate tensor to modulate the addition
+
+    Returns:
+        x + y if gate is None, otherwise x + y * gate
+    """
+    if x is None and y is None:
+        return None
+    if x is None or y is None:
+        return x if x is not None else y
+    if gate is None:
+        return x + y
+    return x + y * gate
+
+
+modeling_gemma._gated_residual = _gated_residual
+
+
+class PatchedGemmaRMSNorm(nn.Module):
+    """RMS normalization with optional adaptive support (ADARMS)."""
+
+    def __init__(self, dim: int, eps: float = 1e-6, cond_dim: Optional[int] = None):
+        """Initializes the PatchedGemmaRMSNorm.
+
+        Args:
+            dim: The dimension of the input tensor.
+            eps: The epsilon value for numerical stability.
+            cond_dim: The dimension of the conditioning vector (if using ADARMS).
+        """
+        super().__init__()
+        self.eps = eps
+        self.dim = dim
+        self.cond_dim = cond_dim
+
+        # Dense layer for adaptive normalization (if cond_dim is provided)
+        if cond_dim is not None:
+            self.dense = nn.Linear(cond_dim, dim * 3, bias=True)
+            # Initialize with zeros (matches source implementation)
+            nn.init.zeros_(self.dense.weight)
+        else:
+            self.weight = nn.Parameter(torch.zeros(dim))
+            self.dense = None
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        """Applies RMS normalization.
+
+        Args:
+            x: The input tensor.
+
+        Returns:
+            The normalized tensor.
+        """
+        # Compute variance in float32 (like the source implementation)
+        var = torch.mean(torch.square(x.float()), dim=-1, keepdim=True)
+        # Compute normalization in float32
+        normed_inputs = x * torch.rsqrt(var + self.eps)
+        return normed_inputs
+
+    def forward(
+        self, x: torch.Tensor, cond: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Forward pass of the normalization layer.
+
+        Args:
+            x: The input tensor.
+            cond: The conditioning tensor for adaptive normalization.
+
+        Returns:
+            A tuple containing the normalized tensor and the gate tensor (if applicable).
+            If cond is None, the gate tensor will be None.
+
+        Raises:
+            ValueError: If cond dimension does not match the configured cond_dim.
+        """
+        dtype = x.dtype  # original dtype, could be half-precision
+        normed_inputs = self._norm(x)
+
+        if cond is None or self.dense is None:
+            # regular RMSNorm
+            # scale by learned parameter in float32 (matches source implementation)
+            normed_inputs = normed_inputs * (1.0 + self.weight.float())
+            return normed_inputs.to(dtype), None  # return in original dtype with None gate
+
+        # adaptive RMSNorm (if cond is provided and dense layer exists)
+        if cond.shape[-1] != self.cond_dim:
+            raise ValueError(f"Expected cond dimension {self.cond_dim}, got {cond.shape[-1]}")
+
+        modulation = self.dense(cond)
+        # Reshape modulation to broadcast properly: [batch, 1, features] for [batch, seq, features]
+        if len(x.shape) == 3:  # [batch, seq, features]
+            modulation = modulation.unsqueeze(1)
+
+        scale, shift, gate = torch.chunk(modulation, 3, dim=-1)
+
+        normed_inputs = normed_inputs * (1 + scale.to(torch.float32)) + shift.to(torch.float32)
+
+        return normed_inputs.to(dtype), gate.to(dtype)
+
+    def extra_repr(self) -> str:
+        """Returns the extra representation of the module."""
+        if hasattr(self, "weight") and self.weight is not None:
+            repr_str = f"{tuple(self.weight.shape)}, eps={self.eps}"
+        else:
+            repr_str = f"dim={self.dim}, eps={self.eps}"
+        if self.dense is not None:
+            repr_str += f", adaptive=True, cond_dim={self.cond_dim}"
+        return repr_str
+
+
+# Apply patches
+modeling_gemma.GemmaRMSNorm = PatchedGemmaRMSNorm
+
+
+def patched_gemma_decoder_layer_init(self, config: GemmaConfig, layer_idx: int):
+    """Initializes a GemmaDecoderLayer with potential ADARMS support.
+
+    Args:
+        self: The GemmaDecoderLayer instance.
+        config: The configuration object.
+        layer_idx: The index of the layer.
+    """
+    modeling_gemma.GradientCheckpointingLayer.__init__(self)
+    self.hidden_size = config.hidden_size
+
+    self.self_attn = modeling_gemma.GemmaAttention(config=config, layer_idx=layer_idx)
+
+    self.mlp = modeling_gemma.GemmaMLP(config)
+
+    cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None
+    self.input_layernorm = modeling_gemma.GemmaRMSNorm(
+        config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim
+    )
+    self.post_attention_layernorm = modeling_gemma.GemmaRMSNorm(
+        config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim
+    )
+
+
+modeling_gemma.GemmaDecoderLayer.__init__ = patched_gemma_decoder_layer_init
+
+
+def patched_gemma_model_init(self, config: GemmaConfig):
+    """Initializes the GemmaModel with potential ADARMS support.
+
+    Args:
+        self: The GemmaModel instance.
+        config: The configuration object.
+    """
+    modeling_gemma.GemmaPreTrainedModel.__init__(self, config)
+    self.padding_idx = config.pad_token_id
+    self.vocab_size = config.vocab_size
+
+    self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+    self.layers = nn.ModuleList(
+        [modeling_gemma.GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+    )
+
+    cond_dim = getattr(config, "adarms_cond_dim", None) if getattr(config, "use_adarms", False) else None
+    self.norm = modeling_gemma.GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps, cond_dim=cond_dim)
+    self.rotary_emb = modeling_gemma.GemmaRotaryEmbedding(config=config)
+    self.gradient_checkpointing = False
+
+    # Initialize weights and apply final processing
+    self.post_init()
+
+
+modeling_gemma.GemmaModel.__init__ = patched_gemma_model_init
+
+
+def patched_gemma_pretrained_model_init_weights(self, module: nn.Module):
+    """Initializes the weights of the GemmaPreTrainedModel.
+
+    Args:
+        self: The GemmaPreTrainedModel instance.
+        module: The module to initialize.
+    """
+    std = self.config.initializer_range
+    if isinstance(module, nn.Linear):
+        module.weight.data.normal_(mean=0.0, std=std)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    elif isinstance(module, nn.Embedding):
+        module.weight.data.normal_(mean=0.0, std=std)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    elif isinstance(module, modeling_gemma.GemmaRMSNorm):
+        if hasattr(module, "weight"):
+            module.weight.data.fill_(1.0)
+
+
+modeling_gemma.GemmaPreTrainedModel._init_weights = patched_gemma_pretrained_model_init_weights
+
+
+def patched_paligemma_model_get_image_features(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+    """Obtains image last hidden states from the vision tower and apply multimodal projection.
 
     Args:
-        distribution_name: The name of the distribution to retrieve metadata for.
+        self: The PaliGemmaModel instance.
+        pixel_values: The tensors corresponding to the input images.
+            Shape: (batch_size, channels, height, width).
 
     Returns:
-        The distribution metadata object.
+        Image feature tensor of shape (num_images, image_length, embed_dim).
     """
-    if distribution_name == "transformers":
-        distribution_name = "opentau-transformers"
-    return _orig_distribution(distribution_name)
+    image_outputs = self.vision_tower(pixel_values)
+    selected_image_feature = image_outputs.last_hidden_state
+    image_features = self.multi_modal_projector(selected_image_feature)
+    return image_features
 
 
-# Apply the patch
-importlib.metadata.distribution = distribution
+PaliGemmaModel.get_image_features = patched_paligemma_model_get_image_features

{opentau-0.1.1.dist-info → opentau-0.2.0.dist-info}/METADATA

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: opentau
-Version: 0.1.1
+Version: 0.2.0
 Summary: OpenTau: Tensor's VLA Training Infrastructure for Real-World Robotics in Pytorch
-Author-email: Shuheng Liu <wish1104@icloud.com>, William Yue <williamyue37@gmail.com>, Akshay Shah <akshayhitendrashah@gmail.com>, Xingrui Gu <xingrui_gu@berkeley.edu>
+Author-email: Shuheng Liu <wish1104@icloud.com>, William Yue <williamyue37@gmail.com>, Akshay Shah <akshayhitendrashah@gmail.com>
 License: Apache-2.0
 Project-URL: homepage, https://github.com/TensorAuto/OpenTau
 Project-URL: issues, https://github.com/TensorAuto/OpenTau/issues
@@ -41,9 +41,9 @@ Requires-Dist: pynput>=1.7.7
 Requires-Dist: pyzmq>=26.2.1
 Requires-Dist: rerun-sdk>=0.21.0
 Requires-Dist: termcolor>=2.4.0
-Requires-Dist: torch<2.8.0,>=2.7.1
+Requires-Dist: torch>=2.7.1
 Requires-Dist: torchcodec<0.5.0,>=0.4.0; sys_platform != "win32" and (sys_platform != "linux" or (platform_machine != "aarch64" and platform_machine != "arm64" and platform_machine != "armv7l")) and (sys_platform != "darwin" or platform_machine != "x86_64")
-Requires-Dist: torchvision<0.23.0,>=0.22.1
+Requires-Dist: torchvision>=0.22.1
 Requires-Dist: wandb>=0.16.3
 Requires-Dist: zarr>=2.17.0
 Requires-Dist: scikit-learn>=1.7.1
@@ -52,7 +52,7 @@ Requires-Dist: onnxruntime>=1.22.1; sys_platform == "darwin" or platform_machine
 Requires-Dist: onnxruntime-gpu>=1.22.0; (sys_platform == "linux" and platform_machine == "x86_64") or (sys_platform == "win32" and (platform_machine == "AMD64" or platform_machine == "x86_64"))
 Requires-Dist: onnxscript>=0.3.1
 Requires-Dist: onnx-ir>=0.1.4
-Requires-Dist: opentau-transformers==4.53.3
+Requires-Dist: transformers==4.53.3
 Requires-Dist: scipy>=1.15.2
 Requires-Dist: pytest>=8.1.0
 Requires-Dist: pytest-cov>=5.0.0
@@ -62,6 +62,10 @@ Requires-Dist: scikit-image>=0.23.2
 Requires-Dist: pandas>=2.2.2
 Requires-Dist: accelerate>=1.4.0
 Requires-Dist: deepspeed>=0.17.1
+Requires-Dist: gymnasium[other]>=0.29
+Requires-Dist: grpcio>=1.60.0
+Requires-Dist: grpcio-tools>=1.60.0
+Requires-Dist: protobuf>=4.25.0
 Provides-Extra: dev
 Requires-Dist: pre-commit>=3.7.0; extra == "dev"
 Requires-Dist: debugpy>=1.8.1; extra == "dev"
@@ -93,10 +97,11 @@ Requires-Dist: libero; extra == "libero"
 Requires-Dist: numpy<2; extra == "libero"
 Requires-Dist: gym<0.27,>=0.25; extra == "libero"
 Requires-Dist: pyopengl-accelerate==3.1.7; sys_platform == "linux" and extra == "libero"
-Requires-Dist: gymnasium[other]>=0.29; extra == "libero"
 Requires-Dist: mujoco>=3.1.6; sys_platform == "linux" and extra == "libero"
 Requires-Dist: pyopengl==3.1.7; sys_platform == "linux" and extra == "libero"
 Requires-Dist: numpy==1.26.4; sys_platform == "linux" and extra == "libero"
+Provides-Extra: urdf
+Requires-Dist: rerun-sdk>=0.28.2; extra == "urdf"
 Dynamic: license-file
 
 <p align="center">
@@ -105,6 +110,19 @@ Dynamic: license-file
   </a>
 </p>
 
+<p align="center">
+  <a href="https://github.com/TensorAuto/OpenTau/actions/workflows/cpu_test.yml?query=branch%3Amain"><img src="https://github.com/TensorAuto/OpenTau/actions/workflows/cpu_test.yml/badge.svg?branch=main" alt="CPU Tests"></a>
+  <a href="https://github.com/TensorAuto/OpenTau/actions/workflows/gpu_test.yml"><img src="https://github.com/TensorAuto/OpenTau/actions/workflows/gpu_test.yml/badge.svg" alt="Nightly GPU Tests"></a>
+  <a href="https://github.com/TensorAuto/OpenTau/actions/workflows/regression_test.yml"><img src="https://github.com/TensorAuto/OpenTau/actions/workflows/regression_test.yml/badge.svg" alt="Nightly Regression Tests"></a>
+  <a href="https://opentau.readthedocs.io/en/latest/?badge=latest"><img src="https://readthedocs.org/projects/opentau/badge/?version=latest" alt="Documentation"></a>
+  <a href="https://pypi.org/project/opentau/"><img src="https://img.shields.io/pypi/v/opentau" alt="Version"></a>
+  <a href="https://pypi.org/project/opentau/"><img src="https://img.shields.io/pypi/status/opentau" alt="Status"></a>
+  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/pypi/pyversions/opentau" alt="Python versions"></a>
+  <a href="https://github.com/TensorAuto/OpenTau/blob/main/LICENSE"><img src="https://img.shields.io/badge/License-Apache%202.0-blue.svg" alt="License"></a>
+  <a href="https://hub.docker.com/r/tensorauto/opentau"><img src="https://img.shields.io/docker/v/tensorauto/opentau?label=Docker" alt="Docker"></a>
+  <a href="https://github.com/pre-commit/pre-commit"><img src="https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit" alt="pre-commit"></a>
+</p>
+
 # OpenTau - Train VLA models with state-of-the-art techniques by Tensor
 
 At Tensor, we are pushing the frontier of large foundation models for physical AI. In robot learning, a vision-language-action (VLA) model is a multimodal foundation model that integrates vision, language, and action. Today, VLA represents the leading approach for embodied AI, spanning autonomous driving, robot manipulation, and navigation.
@@ -122,17 +140,19 @@ Whether you use the official OpenPi codebase or LeRobot’s reimplementation, yo
 
 OpenTau ($\tau$) is a tool developed by *[Tensor][1]* to bridge this gap, and we also use it internally to train our proprietary in-house models. Our goal is to help you train VLAs on any dataset while fully leveraging state-of-the-art techniques. We plan to continuously upgrade this repository to keep pace with the state of the art in the robotics community.
 
-| Features                                                 | OpenPi                  | LeRobot                          | **OpenTau** |
-| -------------------------------------------------------: | :---------------------: | :------------------------------: | :---------: |
-| Co-training with Heterogeneous Datasets                  | ❌                       | ❌                                | ✅           |
-| Discrete Actions Training in $\pi_{0.5}$                 | ❌                       | ❌                                | ✅           |
-| Knowledge Insulation (KI) between VLM and Action Decoder | ❌                       | ❌                                | ✅           |
-| Dropout Layers in PaliGemma                              | ✅ (Jax) <br>❌ (PyTorch) | ❌                                | ✅           |
-| Multi-Node and Multi-GPU Training                        | ❌                       | ✅                                | ✅           |
-| Fully Functioning $\pi_{0.5}$ Checkpoint                 | ✅                       | ❌ <br> (Missing Text Embeddings) | ✅           |
-| Simulation Environments for Evaluating Models            | ❌                       | ✅                                | ✅           |
-| $\pi^{*}_{0.6}$ style Reinforcement Learning Pipeline    | ❌                       | ❌                                | ✅           |
-| Framework                                                | Jax / PyTorch           | PyTorch                          | PyTorch     |
+|                                                 Features |         OpenPi          |             LeRobot              | **OpenTau** |
+|---------------------------------------------------------:|:-----------------------:|:--------------------------------:|:-----------:|
+|                  Co-training with Heterogeneous Datasets |            ❌            |                ❌                 |      ✅      |
+|                 Discrete Actions Training in $\pi_{0.5}$ |            ❌            |                ❌                 |      ✅      |
+| Knowledge Insulation (KI) between VLM and Action Decoder |            ❌            |                ❌                 |      ✅      |
+|                              Dropout Layers in PaliGemma | ✅ (Jax) <br>❌ (PyTorch) |                ❌                 |      ✅      |
+|                        Multi-Node and Multi-GPU Training |            ❌            |                ✅                 |      ✅      |
+|                 Fully Functioning $\pi_{0.5}$ Checkpoint |            ✅            | ❌ <br> (Missing Text Embeddings) |      ✅      |
+|                       Visualize dataset with URDF models |            ❌            |                ❌                 |      ✅      |
+|            Simulation Environments for Evaluating Models |            ❌            |                ✅                 |      ✅      |
+|                 Create Validation Splits During Training |            ❌            |                ❌                 |      ✅      |
+|    $\pi^{*}_{0.6}$ style Reinforcement Learning Pipeline |            ❌            |                ❌                 |      ✅      |
+|                                                Framework |      Jax / PyTorch      |             PyTorch              |   PyTorch   |
 
 ## Quick Start
 If you are familiar with LeRobot, getting started with OpenTau is very easy.

{opentau-0.1.1.dist-info → opentau-0.2.0.dist-info}/RECORD

@@ -2,26 +2,27 @@ opentau/__init__.py,sha256=fIqOYsZsF-bitUI-4taSNke_D1YJYCehGseZNe29GG0,6756
 opentau/__version__.py,sha256=junxoss59Jz_hmg3YnzhpVk_Q5Fo6uha23P1ET81N1c,889
 opentau/constants.py,sha256=-_CbJujCp6hbBjJHgYMguCTcSAkVkmdpM4wHqZp7vRQ,2020
 opentau/configs/__init__.py,sha256=hC-KkeCfq1mtMw9WjPCZfOTxrzQWW7hAa1w8BRC_Bqw,784
-opentau/configs/default.py,sha256=MEoZyzK8olVXVHDy3FHQKK43-ULq2UQv3almymbmdCI,14387
+opentau/configs/default.py,sha256=T5QS84VNcoFaek8Ry7BsDAvsNsV5RLu2L2y2VvO3THo,15273
+opentau/configs/deployment.py,sha256=kem4WCkAKV6H_JUFEVyzw4h-EhQTPa6-9pY65Wr9wDQ,3106
 opentau/configs/libero.py,sha256=CrRfiCBYOw7hVqv6orH_ahNyQudj6iyqHtZM9YpdvzE,4688
 opentau/configs/parser.py,sha256=Pb7sw6yx38F31Aqw1J7wK6BRzfBA7DutywShyP_t9bY,14890
 opentau/configs/policies.py,sha256=06oUJx0B4V6krRwyjH1goTYM3RIpRozg4SSwcVJurG4,11667
 opentau/configs/reward.py,sha256=t7S8_RpEy31fAP4v_ygB-ETvaUR6OyrXmS1JNSz3cOk,1537
-opentau/configs/train.py,sha256=a-c-s2zpCUkK8n0DqBPGGaJ87UcVNF02RcLa1pH843Y,18049
+opentau/configs/train.py,sha256=nn2QX151wI-R-qbghyMkSv1miSPvNSUtsX2Swu-HVGU,18396
 opentau/configs/types.py,sha256=DvKasR2v0ecSmozL0YD4S-64OeuDYhVBhtspxUDV5u0,2453
 opentau/datasets/__init__.py,sha256=oLfV9vfFOg7o2XIRFiN5zOf529FafIkPwqFG7iUX4gc,4248
 opentau/datasets/backward_compatibility.py,sha256=ENVQk9QDPCip6NfAxNF6Vo4VvyCWb0acV-ZxcJBsB6o,3459
 opentau/datasets/compute_stats.py,sha256=N359TDuJicLKMtxxy0JVEcUtnTOB57gL5G8e9Dq0gMQ,13069
 opentau/datasets/dataset_mixture.py,sha256=8UWjY9oKn9jEMe-e9Dy6no1p_21H0kXKv8A10Ku_8_o,19850
-opentau/datasets/factory.py,sha256=NKWpbuNBve0PsmK1midj8g1IpQapeHn-VrxCOC3X4eI,10480
+opentau/datasets/factory.py,sha256=KVN8XEjeIdfTMohyftghG3dsM50UPjv05lL3eS5aTI4,12116
 opentau/datasets/image_writer.py,sha256=JYCkImHFYpLuE88t16cYqXqQS7EHS7g6kLWXPCJmWgw,11072
-opentau/datasets/lerobot_dataset.py,sha256=rz_3BcXqpcIzYr0NEVmkfLf1dY7vcTdo6zuV1CZkIuI,84747
+opentau/datasets/lerobot_dataset.py,sha256=f15Sy3jWOnuPiXiqB8pGdHqv3MOBZgToJyjcA7ry0JU,84778
 opentau/datasets/online_buffer.py,sha256=x14P8tBz25s-hRlE8loFJs5CAvh65RGWeogF271hiF0,19671
 opentau/datasets/sampler.py,sha256=5g-6prsWItVjqkt1J7mA9JPNQPhDSFx3r6rA4JphP9U,4012
 opentau/datasets/standard_data_format_mapping.py,sha256=wEKilksMUjJGeIhvyLuR9qhyhtiJMK1e1AzCkbyx-l4,9667
 opentau/datasets/transforms.py,sha256=pr_8vOEDUoWu7aOUdnI0_wgetsFuie3I2UYFrcStG1k,12976
 opentau/datasets/utils.py,sha256=bZ0Q8KPZMWe9fLdrqJqslgDxI9sa8uxqPQTxEyWwDKw,45062
-opentau/datasets/video_utils.py,sha256=NY20Et6SKWLdG4EjTNdXhpPqWEFON-UccIn_P2YukSQ,21810
+opentau/datasets/video_utils.py,sha256=AUNUKr4IrDVetfYjzZj9Uq4GKrdHrzVYTcMoF1Jlggw,21968
 opentau/datasets/grounding/__init__.py,sha256=ojvdscCIjp5EQxptFAjPgvjKGZa_Xk9LLZ2wNUebWFw,3139
 opentau/datasets/grounding/base.py,sha256=FDAn2QPQHNBB7mzD7IQ2Bz882Dt1RPasBTgskXqKbP4,5773
 opentau/datasets/grounding/clevr.py,sha256=lNZ0hr5EhQKTh-eg35zujybcAo5p8JQEn9dW8DJhOjI,3983
@@ -59,9 +60,9 @@ opentau/policies/pi0/configuration_pi0.py,sha256=94EG2QlraDsPjD0zyuGwKPqToqV_ayP
 opentau/policies/pi0/modeling_pi0.py,sha256=rz1S7hDOVEv12sN0ECGupddKwQVXMqdvm7G_OooWZLA,37442
 opentau/policies/pi0/paligemma_with_expert.py,sha256=j9P6SL7MVP11MgyJqNnsrZAlOctWmqwDaJJAx4z9F84,20724
 opentau/policies/pi05/__init__.py,sha256=VcIjZwlRW1JChRHqAK9Vz4JAIEP40RrP-W-UdyR6xk4,821
-opentau/policies/pi05/configuration_pi05.py,sha256=ucgCC3BaIC6rcnotMYbElTq7ymPZh4xDGghxCseK33M,9307
-opentau/policies/pi05/modeling_pi05.py,sha256=sF4OPGQWQ0eJevGDRZDPUItBkBC5Wp6WdpwsN4uqS14,50639
-opentau/policies/pi05/paligemma_with_expert.py,sha256=nxBfUwBt6S4WwPDkn8LW3lHFcOib4CtoZzJEq4VKlck,23328
+opentau/policies/pi05/configuration_pi05.py,sha256=GmENtmgvI5q3gQQlZnH6RalV5msU5gwtTK_imbNH6a8,9367
+opentau/policies/pi05/modeling_pi05.py,sha256=sBmnP2cqaVHvCGGpg624LYbun52hZs63UCeadiE4dH4,63327
+opentau/policies/pi05/paligemma_with_expert.py,sha256=jyYkcOVMxMNYJtmbR3qZPmEoinOY1myqnq1JKmledkc,23407
 opentau/policies/value/__init__.py,sha256=wUP5vdpsvfnREqt4enfwaakzJ-ynX9sLYN14t2zEtpA,772
 opentau/policies/value/configuration_value.py,sha256=ApjrNKHxvjNlSZ71-BPvanNAJh9GzAK9W0JCiz3mMHs,5951
 opentau/policies/value/modeling_value.py,sha256=21x2EVGFlJbLasJDGTs3b5YMrrChrkuGxrYP7wLjkCY,18594
@@ -78,14 +79,16 @@ opentau/scripts/fake_tensor_training.py,sha256=y4F3CFs2jjpIJcT1wKvsrgFEebU9QFzba
 opentau/scripts/get_advantage_and_percentiles.py,sha256=JdjlADYzdS1Jc_19H6lLYMRnPlWxeckRSUQqwqb0rC4,8993
 opentau/scripts/high_level_planner_inference.py,sha256=nbXr8Hp64YGeprMTpT8kvT_NgpBlI02CUlO6Mm2Js_E,3846
 opentau/scripts/inference.py,sha256=_lp9YjPzarAnjiA8k2jBlIKZxza6PEHw--UyaqLPdNo,2110
-opentau/scripts/launch.py,sha256=cHhE_LuPnmkSkm6F-5XEgwmPDmbhVnLTVWDZ0lqfjPo,2663
-opentau/scripts/libero_simulation_parallel.py,sha256=qMee6T0EwMoAT1J2u8X4w8rsbOJYwyqD3LRAPe2Ta1g,13105
-opentau/scripts/libero_simulation_sequential.py,sha256=xFSUQEuyai20QD-pYitp-UJPGE9zlaaIu4YSO0bhYKg,4775
+opentau/scripts/launch.py,sha256=L_KlkcJpcOsSMGlBKSmtTUyzb7q8tH4FkmuHx8lEdDI,2845
 opentau/scripts/nav_high_level_planner_inference.py,sha256=z2WHw68NWi-fJUd5TV4CrJHzxo-L7e2UliGjfOlqifM,1878
-opentau/scripts/train.py,sha256=nkvsvna5yliphp7pwVyFY6yBwCA_kmffyohRO2wjiHU,16850
-opentau/scripts/visualize_dataset.py,sha256=_xGfAXQqhjGYMi__6L7qRH2xS5XQ2-GQRXjNw3KXMlY,10109
-opentau/scripts/visualize_dataset_html.py,sha256=gEX-E5fFqBhINthf7xLMICHySvw9e3Kcf1HPRnJIyug,17979
+opentau/scripts/train.py,sha256=dp_366gKpFeIcv2tfDkuuFGCdPyP74lkENETZpwR5m4,20547
+opentau/scripts/visualize_dataset.py,sha256=ZfB7Qbsl3RGqu8k7n6CK6iRbhuYhYSX167tta2b01NQ,13625
 opentau/scripts/zero_to_fp32.py,sha256=Rkl1ZczytKix9vGMg0EELzdJYFqUM1yB9p3xvSaK9k8,33272
+opentau/scripts/grpc/__init__.py,sha256=wBiZyRqF1PCpZHgqwHjSZaaeFRHLOX4ZggCXbAzngOs,799
+opentau/scripts/grpc/client.py,sha256=PbuAb14izNAItspdvppItpv6gaycBs2_mdcvrdtAXnQ,20181
+opentau/scripts/grpc/robot_inference_pb2.py,sha256=Se7elLBIHeoi3UIBXrn02w8rZVOg-gpBcuirt23d8Tg,4125
+opentau/scripts/grpc/robot_inference_pb2_grpc.py,sha256=LymJbsaei6C-pp9ffz2ETCVIhq4SRh--LdK8ne0-yug,7659
+opentau/scripts/grpc/server.py,sha256=x6BA0F0uYIXeBw4mnOXaYb2-y8zEbgMXoiRdtNvyq1g,11146
 opentau/utils/__init__.py,sha256=hIUeGPpZHf2AVf0-5C2p0BOcY0cFHCTT5yHn-SpEPwY,856
 opentau/utils/accelerate_utils.py,sha256=vXnSGo1hXCUNof-oNKLMJ_SOMjpKhpZ1gx21ObSsopI,2630
 opentau/utils/benchmark.py,sha256=jVli6gdBRMXAqNM3AIi43a0N_O1CLQMbKXsPK_e2y3s,3063
@@ -99,11 +102,11 @@ opentau/utils/logging_utils.py,sha256=zd7ypmk7aqVposPhA7Kg-PYrstapY4MsuTklsTD4r4
 opentau/utils/monkey_patch.py,sha256=cVgZ1N-NNVnlRKPA1dwO9FM4IbxR0V_Hbil6p-6knhA,9558
 opentau/utils/random_utils.py,sha256=k3Ab3Y98LozGdsBzKoP8xSsFTcnaRqUzY34BsETCrrA,9102
 opentau/utils/train_utils.py,sha256=0d7yvk8wlP-75pwB55gr095b_b1sWG5nlqdVxyH6_o0,6796
-opentau/utils/transformers_patch.py,sha256=3r51KjUQIQQLpG3RQHgA7OP4kboHsH2y5upFBCMgiKI,1781
+opentau/utils/transformers_patch.py,sha256=rPG2Yn7GQr2gCEykhW42uOoKP_jdAMx4p3q-IUcGYDI,10045
 opentau/utils/utils.py,sha256=DrMStfjBEkw_8WVhYMnCQJNBxMeozIJ8LBSpOtMQhFM,15760
-opentau-0.1.1.dist-info/licenses/LICENSE,sha256=tl3_NkxplsgU86xSvEWnDlE1UR_JsIvGo7t4hPtsIbE,27680
-opentau-0.1.1.dist-info/METADATA,sha256=CZpUmmtkXZ9K0dNMU3F4K_xwDMwa3mo2hFq16N23PnM,10951
-opentau-0.1.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-opentau-0.1.1.dist-info/entry_points.txt,sha256=CmA0PWn2JepmMtBY4RDfZ48KaMCYVYF81Xjk55Jd9HM,153
-opentau-0.1.1.dist-info/top_level.txt,sha256=7_yrS4x5KSeTRr2LICTCNOZmF-1_kSOFPKHvtJPL1Dw,8
-opentau-0.1.1.dist-info/RECORD,,
+opentau-0.2.0.dist-info/licenses/LICENSE,sha256=tl3_NkxplsgU86xSvEWnDlE1UR_JsIvGo7t4hPtsIbE,27680
+opentau-0.2.0.dist-info/METADATA,sha256=OUN-KCCeTA1cFdlzmwRKecxepcmjpie6-AbbIkoELhI,12992
+opentau-0.2.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+opentau-0.2.0.dist-info/entry_points.txt,sha256=NGF_MWpSKri0lvjR9WGN4pBUap8B-z21f7XMluxc1M4,208
+opentau-0.2.0.dist-info/top_level.txt,sha256=7_yrS4x5KSeTRr2LICTCNOZmF-1_kSOFPKHvtJPL1Dw,8
+opentau-0.2.0.dist-info/RECORD,,

{opentau-0.1.1.dist-info → opentau-0.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

{opentau-0.1.1.dist-info → opentau-0.2.0.dist-info}/entry_points.txt

@@ -1,4 +1,5 @@
 [console_scripts]
+opentau-dataset-viz = opentau.scripts.launch:visualize
 opentau-eval = opentau.scripts.launch:eval
 opentau-export = opentau.scripts.launch:export
 opentau-train = opentau.scripts.launch:train