opentau 0.1.1__tar.gz → 0.2.0__tar.gz

{opentau-0.1.1/src/opentau.egg-info → opentau-0.2.0}/PKG-INFO

@@ -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 → opentau-0.2.0}/README.md

@@ -4,6 +4,19 @@
   </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.
@@ -21,17 +34,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 → opentau-0.2.0}/pyproject.toml

@@ -20,13 +20,12 @@ huggingface = "https://huggingface.co/TensorAuto"
 
 [project]
 name = "opentau"
-version = "0.1.1"
+version = "0.2.0"
 description = "OpenTau: Tensor's VLA Training Infrastructure for Real-World Robotics in Pytorch"
 authors = [
     { name = "Shuheng Liu", email = "wish1104@icloud.com" },
     { name = "William Yue", email = "williamyue37@gmail.com" },
     { name = "Akshay Shah", email = "akshayhitendrashah@gmail.com" },
-    { name = "Xingrui Gu", email = "xingrui_gu@berkeley.edu" }
 ]
 readme = "README.md"
 license = { text = "Apache-2.0" }
@@ -65,9 +64,9 @@ dependencies = [
     "pyzmq>=26.2.1",
     "rerun-sdk>=0.21.0",
     "termcolor>=2.4.0",
-    "torch>=2.7.1,<2.8.0",
+    "torch>=2.7.1",
     "torchcodec>=0.4.0, <0.5.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')",
-    "torchvision>=0.22.1,<0.23.0",
+    "torchvision>=0.22.1",
     "wandb>=0.16.3",
     "zarr>=2.17.0",
     "scikit-learn>=1.7.1",
@@ -76,7 +75,7 @@ dependencies = [
     "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'))) ",
     "onnxscript>=0.3.1",
     "onnx-ir>=0.1.4",
-    "opentau-transformers==4.53.3",
+    "transformers==4.53.3",
     "scipy>=1.15.2",
     "pytest>=8.1.0",
     "pytest-cov>=5.0.0",
@@ -85,13 +84,18 @@ dependencies = [
     "scikit-image>=0.23.2",
     "pandas>=2.2.2",
     "accelerate>=1.4.0",
-    "deepspeed>=0.17.1"
+    "deepspeed>=0.17.1",
+    "gymnasium[other]>=0.29",
+    "grpcio>=1.60.0",
+    "grpcio-tools>=1.60.0",
+    "protobuf>=4.25.0",
 ]
 
 [project.scripts]
 opentau-train = "opentau.scripts.launch:train"
 opentau-eval = "opentau.scripts.launch:eval"
 opentau-export = "opentau.scripts.launch:export"
+opentau-dataset-viz = "opentau.scripts.launch:visualize"
 
 [project.optional-dependencies]
 dev = ["pre-commit>=3.7.0",
@@ -123,15 +127,26 @@ libero = [
     "numpy<2",
     "gym>=0.25,<0.27",
     "pyopengl-accelerate==3.1.7 ; sys_platform == 'linux'",
-    "gymnasium[other]>=0.29",
     "mujoco>=3.1.6 ; sys_platform == 'linux'",
     "pyopengl==3.1.7 ; sys_platform == 'linux'",
     "numpy==1.26.4 ; sys_platform == 'linux'",
 ]
+urdf = [
+    "rerun-sdk>=0.28.2",
+]
 
 [tool.uv.sources]
 libero = { git = "https://github.com/shuheng-liu/LIBERO" , branch = "master" }  # the official libero repo is misconfigured for pip install with git
 
+# libero depends on gym, which depends on numpy 1.x, while rerun only supports urdf in v0.28 which requires numpy 2.x
+[tool.uv]
+conflicts = [
+  [
+    { extra = "libero" },
+    { extra = "urdf" },
+  ],
+]
+
 [tool.setuptools.packages.find]
 where = ["src"]
 
@@ -142,6 +157,10 @@ target-version = "py310"
 [tool.ruff.lint]
 select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"]
 
+[tool.ruff.lint.per-file-ignores]
+# Server must implement gRPC interface with PascalCase method names
+"src/opentau/scripts/grpc/server.py" = ["N802"]
+
 [tool.bandit]
 exclude_dirs = [
     "tests",

{opentau-0.1.1 → opentau-0.2.0}/src/opentau/configs/default.py

@@ -96,6 +96,11 @@ class DatasetConfig:
     data_features_name_mapping: dict[str, str] | None = None
     loss_type_mapping: str | None = None
 
+    # Ratio of the dataset to be used for validation. Please specify a value.
+    # If `val_freq` is set to 0, a validation dataset will not be created and this value will be ignored.
+    # Defaults to 0.05.
+    val_split_ratio: float = 0.05
+
     def __post_init__(self):
         """Validate dataset configuration and register custom mappings if provided."""
         if (self.repo_id is None) == (self.grounding is None):
@@ -148,6 +153,11 @@ class DatasetMixtureConfig:
     image_resample_strategy: str = "nearest"
     # Resample strategy for non-image features, such as action or state
     vector_resample_strategy: str = "nearest"
+    # Ratio of the dataset to be used for validation. Please specify a value.
+    # If `val_freq` is set to 0, a validation dataset will not be created and this value will be ignored.
+    # This value is applied to all datasets in the mixture.
+    # Defaults to 0.05.
+    val_split_ratio: float = 0.05
 
     def __post_init__(self):
         """Validate dataset mixture configuration."""
@@ -163,6 +173,12 @@ class DatasetMixtureConfig:
             raise ValueError(
                 f"`vector_resample_strategy` must be one of ['linear', 'nearest'], got {self.vector_resample_strategy}."
             )
+        if self.val_split_ratio < 0 or self.val_split_ratio > 1:
+            raise ValueError(f"`val_split_ratio` must be between 0 and 1, got {self.val_split_ratio}.")
+
+        # set the val_split_ratio for all datasets in the mixture
+        for dataset_cfg in self.datasets:
+            dataset_cfg.val_split_ratio = self.val_split_ratio
 
 
 @dataclass

opentau-0.2.0/src/opentau/configs/deployment.py

@@ -0,0 +1,85 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Deployment configuration classes for inference servers.
+
+This module provides configuration classes for deploying trained models
+as inference servers, including gRPC server settings.
+"""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ServerConfig:
+    """Configuration for the gRPC inference server.
+
+    This class contains all configuration parameters needed to run a gRPC
+    inference server for robot policy models.
+
+    Args:
+        port: Port number to serve on. Must be between 1 and 65535.
+            Defaults to 50051.
+        max_workers: Maximum number of gRPC worker threads for handling
+            concurrent requests. Defaults to 4.
+        max_send_message_length_mb: Maximum size of outgoing messages in
+            megabytes. Defaults to 100.
+        max_receive_message_length_mb: Maximum size of incoming messages in
+            megabytes. Defaults to 100.
+
+    Raises:
+        ValueError: If port is not in valid range or max_workers is less than 1.
+
+    Example:
+        >>> config = ServerConfig(port=50051, max_workers=8)
+        >>> config.port
+        50051
+    """
+
+    port: int = 50051
+    max_workers: int = 4
+    max_send_message_length_mb: int = 100
+    max_receive_message_length_mb: int = 100
+
+    def __post_init__(self):
+        """Validate server configuration parameters."""
+        if not 1 <= self.port <= 65535:
+            raise ValueError(f"`port` must be between 1 and 65535, got {self.port}.")
+        if self.max_workers < 1:
+            raise ValueError(f"`max_workers` must be at least 1, got {self.max_workers}.")
+        if self.max_send_message_length_mb < 1:
+            raise ValueError(
+                f"`max_send_message_length_mb` must be at least 1, got {self.max_send_message_length_mb}."
+            )
+        if self.max_receive_message_length_mb < 1:
+            raise ValueError(
+                f"`max_receive_message_length_mb` must be at least 1, got {self.max_receive_message_length_mb}."
+            )
+
+    @property
+    def max_send_message_length(self) -> int:
+        """Get maximum send message length in bytes.
+
+        Returns:
+            Maximum send message length in bytes.
+        """
+        return self.max_send_message_length_mb * 1024 * 1024
+
+    @property
+    def max_receive_message_length(self) -> int:
+        """Get maximum receive message length in bytes.
+
+        Returns:
+            Maximum receive message length in bytes.
+        """
+        return self.max_receive_message_length_mb * 1024 * 1024

{opentau-0.1.1 → opentau-0.2.0}/src/opentau/configs/train.py

@@ -32,6 +32,7 @@ from huggingface_hub.errors import HfHubHTTPError
 
 from opentau.configs import parser
 from opentau.configs.default import DatasetMixtureConfig, EvalConfig, WandBConfig
+from opentau.configs.deployment import ServerConfig
 from opentau.configs.policies import PreTrainedConfig
 from opentau.envs.configs import EnvConfig
 from opentau.optim import OptimizerConfig
@@ -116,6 +117,7 @@ class TrainPipelineConfig(HubMixin):
             is disabled. Defaults to 0.
         last_checkpoint_only: If True, only evaluate the last checkpoint.
             Defaults to True.
+        server: Configuration for the gRPC inference server. Defaults to ServerConfig().
     """
 
     dataset_mixture: DatasetMixtureConfig
@@ -163,7 +165,10 @@ class TrainPipelineConfig(HubMixin):
     env: EnvConfig | None = None
     eval: EvalConfig | None = field(default_factory=EvalConfig)
     eval_freq: int = 0  # evaluate every eval_freq steps
+    val_freq: int = 0  # validate every val_freq steps, if 0, then a validation split is not created
     last_checkpoint_only: bool = True
+    # gRPC inference server configuration
+    server: ServerConfig = field(default_factory=ServerConfig)
 
     def __post_init__(self):
         """Initialize post-creation attributes and validate batch size configuration."""

{opentau-0.1.1 → opentau-0.2.0}/src/opentau/datasets/factory.py

@@ -61,7 +61,11 @@ Example:
         >>> dataloader = mixture.get_dataloader()
 """
 
+import copy
+from typing import Tuple, Union
+
 import numpy as np
+import torch
 
 # NOTE: Don't delete; imported for side effects.
 import opentau.datasets.grounding.clevr  # noqa: F401
@@ -151,9 +155,13 @@ def make_dataset(
     cfg: DatasetConfig,
     train_cfg: TrainPipelineConfig,
     return_advantage_input: bool = False,
-) -> BaseDataset:
+) -> Union[BaseDataset, Tuple[BaseDataset, BaseDataset]]:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
+    A train and validation dataset are returned if `train_cfg.val_freq` is greater than 0.
+    The validation dataset is a subset of the train dataset, and is used for evaluation during training.
+    The validation dataset is created by splitting the train dataset into train and validation sets based on `cfg.val_split_ratio`.
+
     Args:
         cfg (DatasetConfig): A DatasetConfig used to create a LeRobotDataset.
         train_cfg (TrainPipelineConfig): A TrainPipelineConfig config which contains a DatasetConfig and a PreTrainedConfig.
@@ -161,10 +169,11 @@ def make_dataset(
             "episode_end_idx", "current_idx", "last_step", "episode_index", and "timestamp". Defaults to False.
 
     Raises:
-        NotImplementedError: The MultiLeRobotDataset is currently deactivated.
+        ValueError: If exactly one of `cfg.grounding` and `cfg.repo_id` is not provided.
+        ValueError: If `cfg.grounding` is not a supported grounding dataset.
 
     Returns:
-        BaseDataset
+        BaseDataset or Tuple[BaseDataset, BaseDataset]: A single dataset or a tuple of (train_dataset, val_dataset) if val_freq > 0.
     """
     image_transforms = ImageTransforms(cfg.image_transforms) if cfg.image_transforms.enable else None
 
@@ -209,12 +218,20 @@ def make_dataset(
                     dataset.meta.stats[key] = {}
                 dataset.meta.stats[key][stats_type] = np.array(stats, dtype=np.float32)
 
+    if train_cfg.val_freq > 0:
+        val_size = int(len(dataset) * cfg.val_split_ratio)
+        train_size = len(dataset) - val_size
+        train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
+        train_dataset.meta = copy.deepcopy(dataset.meta)
+        val_dataset.meta = copy.deepcopy(dataset.meta)
+        return train_dataset, val_dataset
+
     return dataset
 
 
 def make_dataset_mixture(
     cfg: TrainPipelineConfig, return_advantage_input: bool = False
-) -> WeightedDatasetMixture:
+) -> Union[WeightedDatasetMixture, Tuple[WeightedDatasetMixture, WeightedDatasetMixture]]:
     """Creates a dataset mixture from the provided TrainPipelineConfig.
 
     Args:
@@ -223,10 +240,26 @@ def make_dataset_mixture(
             "episode_end_idx", "current_idx", "last_step", "episode_index", and "timestamp". Defaults to False.
 
     Returns:
-        WeightedDatasetMixture: An instance of WeightedDatasetMixture containing the datasets.
+        WeightedDatasetMixture or Tuple[WeightedDatasetMixture, WeightedDatasetMixture]: An instance of WeightedDatasetMixture containing the datasets, or a tuple of (train_mixture, val_mixture) if val_freq > 0.
     """
-    datasets = [
-        make_dataset(dataset_cfg, cfg, return_advantage_input=return_advantage_input)
-        for dataset_cfg in cfg.dataset_mixture.datasets
-    ]
-    return WeightedDatasetMixture(cfg, datasets, cfg.dataset_mixture.weights, cfg.dataset_mixture.action_freq)
+    datasets = []
+    val_datasets = []
+    for dataset_cfg in cfg.dataset_mixture.datasets:
+        res = make_dataset(dataset_cfg, cfg, return_advantage_input=return_advantage_input)
+        if isinstance(res, tuple):
+            datasets.append(res[0])
+            val_datasets.append(res[1])
+        else:
+            datasets.append(res)
+
+    train_mixture = WeightedDatasetMixture(
+        cfg, datasets, cfg.dataset_mixture.weights, cfg.dataset_mixture.action_freq
+    )
+
+    if val_datasets:
+        val_mixture = WeightedDatasetMixture(
+            cfg, val_datasets, cfg.dataset_mixture.weights, cfg.dataset_mixture.action_freq
+        )
+        return train_mixture, val_mixture
+
+    return train_mixture

{opentau-0.1.1 → opentau-0.2.0}/src/opentau/datasets/lerobot_dataset.py

@@ -150,6 +150,7 @@ from opentau.policies.value.configuration_value import ValueConfig
 from opentau.policies.value.reward import (
     calculate_return_bins_with_equal_width,
 )
+from opentau.utils.accelerate_utils import get_proc_accelerator
 from opentau.utils.utils import on_accelerate_main_proc
 
 
@@ -324,8 +325,17 @@ class LeRobotDatasetMetadata(DatasetMetadata):
             if is_valid_version(self.revision):
                 self.revision = get_safe_version(self.repo_id, self.revision)
 
-            (self.root / "meta").mkdir(exist_ok=True, parents=True)
-            self.pull_from_repo(allow_patterns="meta/")
+            # In distributed training, only rank 0 downloads to avoid race conditions
+            # where other ranks read metadata before the download has finished.
+            acc = get_proc_accelerator()
+            if acc is not None and acc.num_processes > 1:
+                if acc.is_main_process:
+                    (self.root / "meta").mkdir(exist_ok=True, parents=True)
+                    self.pull_from_repo(allow_patterns="meta/")
+                acc.wait_for_everyone()
+            else:
+                (self.root / "meta").mkdir(exist_ok=True, parents=True)
+                self.pull_from_repo(allow_patterns="meta/")
             self.load_metadata()
 
     def load_metadata(self) -> None:
@@ -633,7 +643,9 @@ class BaseDataset(torch.utils.data.Dataset):
         For example, {"image_key": torch.zeros(2, 3, 224, 224), "image_key_is_pad": [False, True] } will become
         {
             "image_key": torch.zeros(3, 224, 224),
+            "image_key_local": torch.zeros(3, 224, 224),
             "image_key_is_pad: False,
+            "image_key_local_is_pad": True,
         }.
         """
         raise NotImplementedError
@@ -723,14 +735,6 @@ class BaseDataset(torch.utils.data.Dataset):
             if isinstance(value, torch.Tensor) and value.dtype.is_floating_point:
                 standard_item[key] = value.to(dtype=torch.bfloat16)
 
-        # ensure that non-empty strings contain exactly one newline character at the end of the string
-        for key in ["prompt", "response"]:
-            if standard_item[key].endswith(
-                "\n"
-            ):  # ensure there isn't going to be an extra space at the end after calling replace
-                standard_item[key] = standard_item[key][:-1]
-            standard_item[key] = standard_item[key].replace("\n", " ") + "\n"
-
         return standard_item
 
     def resize_with_pad(self, img, width, height, pad_value=0) -> torch.Tensor:
@@ -1787,16 +1791,12 @@ class LeRobotDataset(BaseDataset):
         cam_keys = {v for k, v in name_map.items() if k.startswith("camera")}
         for k in cam_keys:
             images = item.pop(k)
-            assert len(images) == 2, (
-                f"{k} in {self.__class__} is expected to have length 2, got shape={images.shape}"
-            )
-            item[k + "_local"], item[k] = images
+            if len(images) == 2:
+                item[k + "_local"], item[k] = images
 
-            pads = item.pop(k + "_is_pad")
-            assert len(pads) == 2, (
-                f"{k} in {self.__class__} is expected to have length 2, got shape={pads.shape}"
-            )
-            item[k + "_local_is_pad"], item[k + "_is_pad"] = pads
+            pads = item.get(k + "_is_pad")
+            if hasattr(pads, "__len__") and len(pads) == 2:
+                item[k + "_local_is_pad"], item[k + "_is_pad"] = pads
 
     @staticmethod
     def compute_delta_params(

{opentau-0.1.1 → opentau-0.2.0}/src/opentau/datasets/video_utils.py

@@ -108,6 +108,7 @@ import pyarrow as pa
 import torch
 import torchvision
 from datasets.features.features import register_feature
+from packaging import version
 from PIL import Image
 
 
@@ -117,13 +118,17 @@ def get_safe_default_codec() -> str:
     Returns:
         Backend name: "torchcodec" if available, otherwise "pyav".
     """
-    if importlib.util.find_spec("torchcodec"):
-        return "torchcodec"
-    else:
-        logging.warning(
-            "'torchcodec' is not available in your platform, falling back to 'pyav' as a default decoder"
-        )
+
+    if version.parse(torch.__version__) >= version.parse("2.8.0"):
         return "pyav"
+    else:
+        if importlib.util.find_spec("torchcodec"):
+            return "torchcodec"
+        else:
+            logging.warning(
+                "'torchcodec' is not available in your platform, falling back to 'pyav' as a default decoder"
+            )
+            return "pyav"
 
 
 def decode_video_frames(