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

opentau/scripts/train.py

@@ -17,6 +17,9 @@
 import json
 import logging
 import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
 from contextlib import nullcontext
 from pprint import pformat
 from typing import Any
@@ -176,7 +179,10 @@ def train(cfg: TrainPipelineConfig):
         logging.info("Anomaly detection is disabled.")
 
     logging.info("Creating dataset")
-    dataset = make_dataset_mixture(cfg)
+    if cfg.val_freq > 0:
+        train_dataset, val_dataset = make_dataset_mixture(cfg)
+    else:
+        train_dataset = make_dataset_mixture(cfg)
 
     # Create environment used for evaluating checkpoints during training on simulation data.
     # On real-world data, no need to create an environment as evaluations are done outside train.py,
@@ -188,7 +194,7 @@ def train(cfg: TrainPipelineConfig):
         )
 
     logging.info("Creating policy")
-    policy = make_policy(cfg=cfg.policy, ds_meta=dataset.meta)
+    policy = make_policy(cfg=cfg.policy, ds_meta=train_dataset.meta)
     policy.to(torch.bfloat16)
     logging.info("Creating optimizer and scheduler")
     optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
@@ -203,11 +209,18 @@ def train(cfg: TrainPipelineConfig):
         logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
         logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
-    dataloader = dataset.get_dataloader()
-    policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
-        policy, optimizer, dataloader, lr_scheduler
-    )
-    dl_iter = cycle(dataloader)
+    if cfg.val_freq > 0:
+        train_dataloader = train_dataset.get_dataloader()
+        val_dataloader = val_dataset.get_dataloader()
+        policy, optimizer, train_dataloader, val_dataloader, lr_scheduler = accelerator.prepare(
+            policy, optimizer, train_dataloader, val_dataloader, lr_scheduler
+        )
+    else:
+        train_dataloader = train_dataset.get_dataloader()
+        policy, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
+            policy, optimizer, train_dataloader, lr_scheduler
+        )
+    train_dl_iter = cycle(train_dataloader)
 
     # Register the LR scheduler for checkpointing
     accelerator.register_for_checkpointing(lr_scheduler)
@@ -246,7 +259,7 @@ def train(cfg: TrainPipelineConfig):
         for _ in range(cfg.gradient_accumulation_steps):
             with accelerator.accumulate(policy) if cfg.gradient_accumulation_steps > 1 else nullcontext():
                 logging.debug(f"{step=}, {accelerator.sync_gradients=}")
-                batch = next(dl_iter)
+                batch = next(train_dl_iter)
 
                 train_tracker = update_policy(
                     cfg,
@@ -266,20 +279,21 @@ def train(cfg: TrainPipelineConfig):
         is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
         is_saving_step = (step % cfg.save_freq == 0 or step == cfg.steps) and cfg.save_checkpoint
         is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
+        is_val_step = cfg.val_freq > 0 and step % cfg.val_freq == 0
 
         # Only `train_tracker` on the main process keeps useful statistics,
         #  because we guarded it with if accelerator.is_main_process in the `update_policy` function.
         if is_log_step and accelerator.is_main_process:
             logging.info(train_tracker)
             log_dict = train_tracker.to_dict(use_avg=True)
-            accelerator.log({"Training Loss": log_dict["loss"]}, step=step)
-            accelerator.log({"MSE Loss": log_dict["mse_loss"]}, step=step)
-            accelerator.log({"CE Loss": log_dict["ce_loss"]}, step=step)
-            accelerator.log({"L1 Loss": log_dict["l1_loss"]}, step=step)
-            accelerator.log({"Accuracy": log_dict["accuracy"]}, step=step)
-            accelerator.log({"Learning Rate": log_dict["lr"]}, step=step)
-            accelerator.log({"Grad Norm": log_dict["grad_norm"]}, step=step)
-            accelerator.log({"Num Samples": log_dict["samples"]}, step=step)
+            accelerator.log({"Training/Loss": log_dict["loss"]}, step=step)
+            accelerator.log({"Training/MSE Loss": log_dict["mse_loss"]}, step=step)
+            accelerator.log({"Training/CE Loss": log_dict["ce_loss"]}, step=step)
+            accelerator.log({"Training/L1 Loss": log_dict["l1_loss"]}, step=step)
+            accelerator.log({"Training/Accuracy": log_dict["accuracy"]}, step=step)
+            accelerator.log({"Training/Learning Rate": log_dict["lr"]}, step=step)
+            accelerator.log({"Training/Grad Norm": log_dict["grad_norm"]}, step=step)
+            accelerator.log({"Training/Num Samples": log_dict["samples"]}, step=step)
             train_tracker.reset_averages()
 
         if is_saving_step:
@@ -299,6 +313,70 @@ def train(cfg: TrainPipelineConfig):
                 if cfg.last_checkpoint_only:
                     prune_old_checkpoints(checkpoint_dir)
 
+            accelerator.wait_for_everyone()
+
+        if is_val_step:
+            policy.eval()
+            val_metrics = {
+                "loss": AverageMeter("val_total_loss", ":.3f"),
+                "mse_loss": AverageMeter("val_mse_loss", ":.3f"),
+                "ce_loss": AverageMeter("val_ce_loss", ":.3f"),
+                "l1_loss": AverageMeter("val_l1_loss", ":.3f"),
+                "accuracy": AverageMeter("val_accuracy", ":.3f"),
+            }
+            val_tracker = MetricsTracker(
+                cfg.batch_size * accelerator.num_processes,
+                val_metrics,
+                initial_step=step,
+            )
+
+            logging.info(f"Validation at step {step}...")
+
+            with torch.no_grad():
+                for batch in val_dataloader:
+                    losses = policy.forward(batch)
+                    loss = cfg.loss_weighting["MSE"] * losses["MSE"] + cfg.loss_weighting["CE"] * losses["CE"]
+
+                    # Gather and average metrics across processes
+                    _first_loss_tensor = next(lt for lt in losses.values() if isinstance(lt, torch.Tensor))
+                    zero = torch.tensor(0.0, device=_first_loss_tensor.device, dtype=_first_loss_tensor.dtype)
+
+                    loss = accelerator.gather_for_metrics(loss).mean().item()
+                    mse_loss = (
+                        accelerator.gather_for_metrics(losses["MSE"]).to(dtype=torch.float32).mean().item()
+                    )
+                    ce_loss = (
+                        accelerator.gather_for_metrics(losses["CE"]).to(dtype=torch.float32).mean().item()
+                    )
+                    l1_loss = (
+                        accelerator.gather_for_metrics(losses.get("L1", zero))
+                        .to(dtype=torch.float32)
+                        .mean()
+                        .item()
+                    )
+                    accuracy = (
+                        accelerator.gather_for_metrics(losses.get("Accuracy", zero))
+                        .to(dtype=torch.float32)
+                        .mean()
+                        .item()
+                    )
+
+                    if accelerator.is_main_process:
+                        val_tracker.loss = loss
+                        val_tracker.mse_loss = mse_loss
+                        val_tracker.ce_loss = ce_loss
+                        val_tracker.l1_loss = l1_loss
+                        val_tracker.accuracy = accuracy
+
+            if accelerator.is_main_process:
+                logging.info(val_tracker)
+                val_dict = val_tracker.to_dict(use_avg=True)
+                accelerator.log({"Validation/Loss": val_dict["loss"]}, step=step)
+                accelerator.log({"Validation/MSE Loss": val_dict["mse_loss"]}, step=step)
+                accelerator.log({"Validation/CE Loss": val_dict["ce_loss"]}, step=step)
+                accelerator.log({"Validation/L1 Loss": val_dict["l1_loss"]}, step=step)
+                accelerator.log({"Validation/Accuracy": val_dict["accuracy"]}, step=step)
+
             # This barrier is probably necessary to ensure
             # other processes wait for the main process to finish saving
             accelerator.wait_for_everyone()
@@ -357,7 +435,6 @@ def train(cfg: TrainPipelineConfig):
                 with open(videos_dir / "eval_info.json", "w") as f:
                     json.dump(eval_info, f, indent=2)
 
-        if is_eval_step:
             # This barrier is to ensure all processes finishes evaluation before the next training step
             # Some processes might be slower than others
             accelerator.wait_for_everyone()

opentau/scripts/visualize_dataset.py

@@ -52,7 +52,9 @@ distant$ opentau-dataset-viz --repo-id lerobot/pusht --episode-index 0 --mode di
 import argparse
 import gc
 import logging
+import os
 import time
+import warnings
 from pathlib import Path
 from typing import Iterator
 
@@ -66,6 +68,52 @@ from opentau.configs.default import DatasetMixtureConfig, WandBConfig
 from opentau.configs.train import TrainPipelineConfig
 from opentau.datasets.lerobot_dataset import LeRobotDataset
 
+PERMIT_URDF = hasattr(rr, "urdf")
+if not PERMIT_URDF:
+    warnings.warn(
+        "`rerun.urdf` module not found. Make sure you have rerun >= 0.28.2 installed. "
+        " One way to ensure this is to install OpenTau with the '[urdf]' extra: `pip install opentau[urdf]`.",
+        stacklevel=2,
+    )
+
+
+# Older and newer versions of rerun have different APIs for setting time / sequence
+def _rr_set_sequence(timeline: str, value: int):
+    if hasattr(rr, "set_time_sequence"):
+        rr.set_time_sequence(timeline, value)
+    else:
+        rr.set_time(timeline, sequence=value)
+
+
+def _rr_set_seconds(timeline: str, value: float):
+    if hasattr(rr, "set_time_seconds"):
+        rr.set_time_seconds(timeline, value)
+    else:
+        rr.set_time(timeline, timestamp=value)
+
+
+def _rr_scalar(value: float):
+    """Return a rerun scalar archetype that works across rerun versions.
+
+    Older rerun versions expose `rr.Scalar`, while newer versions expose `rr.Scalars`.
+    This wrapper returns an object suitable for `rr.log(path, ...)` for a single value.
+    """
+    v = float(value)
+
+    # New API (plural archetype)
+    if hasattr(rr, "Scalars"):
+        try:
+            return rr.Scalars(v)
+        except TypeError:
+            # Some versions expect a sequence/array for Scalars.
+            return rr.Scalars([v])
+
+    # Old API
+    if hasattr(rr, "Scalar"):
+        return rr.Scalar(v)
+
+    raise AttributeError("rerun has neither `Scalar` nor `Scalars` - please upgrade `rerun-sdk`.")
+
 
 def create_mock_train_config() -> TrainPipelineConfig:
     """Create a mock TrainPipelineConfig for dataset visualization.
@@ -123,6 +171,7 @@ def visualize_dataset(
     web_port: int = 9090,
     save: bool = False,
     output_dir: Path | None = None,
+    urdf: Path | None = None,
 ) -> Path | None:
     if save:
         assert output_dir is not None, (
@@ -153,6 +202,17 @@ def visualize_dataset(
     # TODO(rcadene): remove `gc.collect` when rerun version 0.16 is out, which includes a fix
     gc.collect()
 
+    if urdf:
+        rr.log_file_from_path(urdf, static=True)
+        urdf_tree = rr.urdf.UrdfTree.from_file_path(urdf)
+        urdf_joints = [jnt for jnt in urdf_tree.joints() if jnt.joint_type != "fixed"]
+        print(
+            "Assuming the dataset state dimensions correspond to URDF joints in order:\n",
+            "\n".join(f"{i:3d}: {jnt.name}" for i, jnt in enumerate(urdf_joints)),
+        )
+    else:
+        urdf_joints = []
+
     if mode == "distant":
         rr.serve_web_viewer(open_browser=False, web_port=web_port)
 
@@ -161,8 +221,8 @@ def visualize_dataset(
     for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
         # iterate over the batch
         for i in range(len(batch["index"])):
-            rr.set_time_sequence("frame_index", batch["frame_index"][i].item())
-            rr.set_time_seconds("timestamp", batch["timestamp"][i].item())
+            _rr_set_sequence("frame_index", batch["frame_index"][i].item())
+            _rr_set_seconds("timestamp", batch["timestamp"][i].item())
 
             # display each camera image
             for key in dataset.meta.camera_keys:
@@ -172,21 +232,27 @@ def visualize_dataset(
             # display each dimension of action space (e.g. actuators command)
             if "action" in batch:
                 for dim_idx, val in enumerate(batch["action"][i]):
-                    rr.log(f"action/{dim_idx}", rr.Scalar(val.item()))
+                    rr.log(f"action/{dim_idx}", _rr_scalar(val.item()))
 
             # display each dimension of observed state space (e.g. agent position in joint space)
             if "observation.state" in batch:
                 for dim_idx, val in enumerate(batch["observation.state"][i]):
-                    rr.log(f"state/{dim_idx}", rr.Scalar(val.item()))
+                    rr.log(f"state/{dim_idx}", _rr_scalar(val.item()))
+                    # Assuming the state dimensions correspond to URDF joints in order.
+                    # TODO(shuheng): allow overriding with a mapping from state dim to joint name.
+                    if dim_idx < len(urdf_joints):
+                        joint = urdf_joints[dim_idx]
+                        transform = joint.compute_transform(float(val))
+                        rr.log("URDF", transform)
 
             if "next.done" in batch:
-                rr.log("next.done", rr.Scalar(batch["next.done"][i].item()))
+                rr.log("next.done", _rr_scalar(batch["next.done"][i].item()))
 
             if "next.reward" in batch:
-                rr.log("next.reward", rr.Scalar(batch["next.reward"][i].item()))
+                rr.log("next.reward", _rr_scalar(batch["next.reward"][i].item()))
 
             if "next.success" in batch:
-                rr.log("next.success", rr.Scalar(batch["next.success"][i].item()))
+                rr.log("next.success", _rr_scalar(batch["next.success"][i].item()))
 
     if mode == "local" and save:
         # save .rrd locally
@@ -272,7 +338,6 @@ def parse_args() -> dict:
             "Visualize the data by running `rerun path/to/file.rrd` on your local machine."
         ),
     )
-
     parser.add_argument(
         "--tolerance-s",
         type=float,
@@ -283,6 +348,22 @@ def parse_args() -> dict:
             "If not given, defaults to 1e-4."
         ),
     )
+    parser.add_argument(
+        "--urdf",
+        type=Path,
+        default=None,
+        help="Path to a URDF file to load and visualize alongside the dataset.",
+    )
+    parser.add_argument(
+        "--urdf-package-dir",
+        type=Path,
+        default=None,
+        help=(
+            "Root directory of the URDF package to resolve package:// paths. "
+            "You can also set the ROS_PACKAGE_PATH environment variable, "
+            "which will be used if this argument is not provided."
+        ),
+    )
 
     args = parser.parse_args()
     return vars(args)
@@ -293,6 +374,12 @@ def main():
     repo_id = kwargs.pop("repo_id")
     root = kwargs.pop("root")
     tolerance_s = kwargs.pop("tolerance_s")
+    urdf_package_dir = kwargs.pop("urdf_package_dir")
+    if urdf_package_dir:
+        os.environ["ROS_PACKAGE_PATH"] = urdf_package_dir.resolve().as_posix()
+
+    if not PERMIT_URDF:
+        kwargs["urdf"] = None
 
     logging.info("Loading dataset")
     dataset = LeRobotDataset(

opentau/utils/transformers_patch.py

@@ -167,7 +167,10 @@ class PatchedGemmaRMSNorm(nn.Module):
 
     def extra_repr(self) -> str:
         """Returns the extra representation of the module."""
-        repr_str = f"{tuple(self.weight.shape)}, eps={self.eps}"
+        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

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

@@ -1,8 +1,8 @@
 Metadata-Version: 2.4
 Name: opentau
-Version: 0.1.2
+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
@@ -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.2.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=c6bGOz75yEJfYkYqlcfszGkap0VBAMBFXrH8fz1P1WQ,84651
+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,13 +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=kcJtdO1WHYxiHSJpJ_y618tbIvBuGXy8FmH5BEEdVdI,2826
-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=RsON_13oqTm7HN14tGnDBIVAJPCW_-EJzpMHeiXxp24,10492
+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
@@ -98,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=-3Fvf-_owtT_QDUkoGfMWO-pxN5xeQikPljtLMn4MRs,9906
+opentau/utils/transformers_patch.py,sha256=rPG2Yn7GQr2gCEykhW42uOoKP_jdAMx4p3q-IUcGYDI,10045
 opentau/utils/utils.py,sha256=DrMStfjBEkw_8WVhYMnCQJNBxMeozIJ8LBSpOtMQhFM,15760
-opentau-0.1.2.dist-info/licenses/LICENSE,sha256=tl3_NkxplsgU86xSvEWnDlE1UR_JsIvGo7t4hPtsIbE,27680
-opentau-0.1.2.dist-info/METADATA,sha256=Up5VRGhf8RVjBA0mBy6xKA21-6R_t51xvGmG-YgC1EQ,10943
-opentau-0.1.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-opentau-0.1.2.dist-info/entry_points.txt,sha256=NGF_MWpSKri0lvjR9WGN4pBUap8B-z21f7XMluxc1M4,208
-opentau-0.1.2.dist-info/top_level.txt,sha256=7_yrS4x5KSeTRr2LICTCNOZmF-1_kSOFPKHvtJPL1Dw,8
-opentau-0.1.2.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.2.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