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

opentau/datasets/lerobot_dataset.py

@@ -633,7 +633,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
@@ -1787,16 +1789,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/scripts/launch.py

@@ -21,6 +21,7 @@ from types import ModuleType
 import opentau.scripts.eval as eval_script
 import opentau.scripts.export_to_onnx as export_script
 import opentau.scripts.train as train_script
+import opentau.scripts.visualize_dataset as visualize_script
 
 
 def launch(script_module: ModuleType, description: str, use_accelerate: bool = True):
@@ -77,3 +78,7 @@ def eval():
 
 def export():
     launch(export_script, "Launch OpenTau ONNX export", use_accelerate=False)
+
+
+def visualize():
+    launch(visualize_script, "Launch OpenTau visualization", use_accelerate=False)

opentau/scripts/visualize_dataset.py

@@ -14,7 +14,7 @@
 # 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.
-""" Visualize data of **all** frames of any episode of a dataset of type LeRobotDataset.
+"""Visualize data of **all** frames of any episode of a dataset of type LeRobotDataset.
 
 Note: The last frame of the episode doesn't always correspond to a final state.
 That's because our datasets are composed of transition from state to state up to
@@ -30,34 +30,21 @@ Examples:
 
 - Visualize data stored on a local machine:
 ```
-local$ python src/opentau/scripts/visualize_dataset.py \
-    --repo-id lerobot/pusht \
-    --episode-index 0
+local$ opentau-dataset-viz --repo-id lerobot/pusht --episode-index 0
 ```
 
 - Visualize data stored on a distant machine with a local viewer:
 ```
-distant$ python src/opentau/scripts/visualize_dataset.py \
-    --repo-id lerobot/pusht \
-    --episode-index 0 \
-    --save 1 \
-    --output-dir path/to/directory
+distant$ opentau-dataset-viz --repo-id lerobot/pusht --episode-index 0 --save 1 --output-dir path/to/directory
 
 local$ scp distant:path/to/directory/lerobot_pusht_episode_0.rrd .
 local$ rerun lerobot_pusht_episode_0.rrd
 ```
 
 - Visualize data stored on a distant machine through streaming:
-(You need to forward the websocket port to the distant machine, with
-`ssh -L 9087:localhost:9087 username@remote-host`)
 ```
-distant$ python src/opentau/scripts/visualize_dataset.py \
-    --repo-id lerobot/pusht \
-    --episode-index 0 \
-    --mode distant \
-    --ws-port 9087
 
-local$ rerun ws://localhost:9087
+distant$ opentau-dataset-viz --repo-id lerobot/pusht --episode-index 0 --mode distant --web-port 9090
 ```
 
 """
@@ -75,8 +62,34 @@ import torch
 import torch.utils.data
 import tqdm
 
+from opentau.configs.default import DatasetMixtureConfig, WandBConfig
+from opentau.configs.train import TrainPipelineConfig
 from opentau.datasets.lerobot_dataset import LeRobotDataset
-from opentau.scripts.visualize_dataset_html import create_mock_train_config
+
+
+def create_mock_train_config() -> TrainPipelineConfig:
+    """Create a mock TrainPipelineConfig for dataset visualization.
+
+    Returns:
+        TrainPipelineConfig: A mock config with default values.
+    """
+    return TrainPipelineConfig(
+        dataset_mixture=DatasetMixtureConfig(),  # Will be set by the dataset
+        resolution=(224, 224),
+        num_cams=2,
+        max_state_dim=32,
+        max_action_dim=32,
+        action_chunk=50,
+        loss_weighting={"MSE": 1, "CE": 1},
+        num_workers=4,
+        batch_size=8,
+        steps=100_000,
+        log_freq=200,
+        save_checkpoint=True,
+        save_freq=20_000,
+        use_policy_training_preset=True,
+        wandb=WandBConfig(),
+    )
 
 
 class EpisodeSampler(torch.utils.data.Sampler):
@@ -108,7 +121,6 @@ def visualize_dataset(
     num_workers: int = 0,
     mode: str = "local",
     web_port: int = 9090,
-    ws_port: int = 9087,
     save: bool = False,
     output_dir: Path | None = None,
 ) -> Path | None:
@@ -142,7 +154,7 @@ def visualize_dataset(
     gc.collect()
 
     if mode == "distant":
-        rr.serve(open_browser=False, web_port=web_port, ws_port=ws_port)
+        rr.serve_web_viewer(open_browser=False, web_port=web_port)
 
     logging.info("Logging to Rerun")
 
@@ -194,7 +206,7 @@ def visualize_dataset(
             print("Ctrl-C received. Exiting.")
 
 
-def main():
+def parse_args() -> dict:
     parser = argparse.ArgumentParser()
 
     parser.add_argument(
@@ -250,12 +262,6 @@ def main():
         default=9090,
         help="Web port for rerun.io when `--mode distant` is set.",
     )
-    parser.add_argument(
-        "--ws-port",
-        type=int,
-        default=9087,
-        help="Web socket port for rerun.io when `--mode distant` is set.",
-    )
     parser.add_argument(
         "--save",
         type=int,
@@ -279,15 +285,25 @@ def main():
     )
 
     args = parser.parse_args()
-    kwargs = vars(args)
+    return vars(args)
+
+
+def main():
+    kwargs = parse_args()
     repo_id = kwargs.pop("repo_id")
     root = kwargs.pop("root")
     tolerance_s = kwargs.pop("tolerance_s")
 
     logging.info("Loading dataset")
-    dataset = LeRobotDataset(create_mock_train_config(), repo_id, root=root, tolerance_s=tolerance_s)
+    dataset = LeRobotDataset(
+        create_mock_train_config(),
+        repo_id,
+        root=root,
+        tolerance_s=tolerance_s,
+        standardize=False,
+    )
 
-    visualize_dataset(dataset, **vars(args))
+    visualize_dataset(dataset, **kwargs)
 
 
 if __name__ == "__main__":

opentau/utils/transformers_patch.py

@@ -12,37 +12,265 @@
 # 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."""
+        repr_str = f"{tuple(self.weight.shape)}, 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.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: opentau
-Version: 0.1.1
+Version: 0.1.2
 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>
 License: Apache-2.0
@@ -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

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

@@ -15,7 +15,7 @@ opentau/datasets/compute_stats.py,sha256=N359TDuJicLKMtxxy0JVEcUtnTOB57gL5G8e9Dq
 opentau/datasets/dataset_mixture.py,sha256=8UWjY9oKn9jEMe-e9Dy6no1p_21H0kXKv8A10Ku_8_o,19850
 opentau/datasets/factory.py,sha256=NKWpbuNBve0PsmK1midj8g1IpQapeHn-VrxCOC3X4eI,10480
 opentau/datasets/image_writer.py,sha256=JYCkImHFYpLuE88t16cYqXqQS7EHS7g6kLWXPCJmWgw,11072
-opentau/datasets/lerobot_dataset.py,sha256=rz_3BcXqpcIzYr0NEVmkfLf1dY7vcTdo6zuV1CZkIuI,84747
+opentau/datasets/lerobot_dataset.py,sha256=c6bGOz75yEJfYkYqlcfszGkap0VBAMBFXrH8fz1P1WQ,84651
 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
@@ -78,13 +78,12 @@ 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/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/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/visualize_dataset.py,sha256=RsON_13oqTm7HN14tGnDBIVAJPCW_-EJzpMHeiXxp24,10492
 opentau/scripts/zero_to_fp32.py,sha256=Rkl1ZczytKix9vGMg0EELzdJYFqUM1yB9p3xvSaK9k8,33272
 opentau/utils/__init__.py,sha256=hIUeGPpZHf2AVf0-5C2p0BOcY0cFHCTT5yHn-SpEPwY,856
 opentau/utils/accelerate_utils.py,sha256=vXnSGo1hXCUNof-oNKLMJ_SOMjpKhpZ1gx21ObSsopI,2630
@@ -99,11 +98,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=-3Fvf-_owtT_QDUkoGfMWO-pxN5xeQikPljtLMn4MRs,9906
 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.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.1.1.dist-info → opentau-0.1.2.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

opentau/scripts/visualize_dataset_html.py

@@ -1,507 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-# 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.
-""" Visualize data of **all** frames of any episode of a dataset of type LeRobotDataset.
-
-Note: The last frame of the episode doesnt always correspond to a final state.
-That's because our datasets are composed of transition from state to state up to
-the antepenultimate state associated to the ultimate action to arrive in the final state.
-However, there might not be a transition from a final state to another state.
-
-Note: This script aims to visualize the data used to train the neural networks.
-~What you see is what you get~. When visualizing image modality, it is often expected to observe
-lossly compression artifacts since these images have been decoded from compressed mp4 videos to
-save disk space. The compression factor applied has been tuned to not affect success rate.
-
-Example of usage:
-
-- Visualize data stored on a local machine:
-```bash
-local$ python src/opentau/scripts/visualize_dataset_html.py \
-    --repo-id lerobot/pusht
-
-local$ open http://localhost:9090
-```
-
-- Visualize data stored on a distant machine with a local viewer:
-```bash
-distant$ python src/opentau/scripts/visualize_dataset_html.py \
-    --repo-id lerobot/pusht
-
-local$ ssh -L 9090:localhost:9090 distant  # create a ssh tunnel
-local$ open http://localhost:9090
-```
-
-- Select episodes to visualize:
-```bash
-python src/opentau/scripts/visualize_dataset_html.py \
-    --repo-id lerobot/pusht \
-    --episodes 7 3 5 1 4
-```
-"""
-
-import argparse
-import csv
-import json
-import logging
-import re
-import shutil
-import tempfile
-from io import StringIO
-from pathlib import Path
-
-import numpy as np
-import pandas as pd
-import requests
-from flask import Flask, redirect, render_template, request, url_for
-
-from opentau import available_datasets
-from opentau.configs.default import DatasetMixtureConfig, WandBConfig
-from opentau.configs.train import TrainPipelineConfig
-from opentau.datasets.lerobot_dataset import LeRobotDataset
-from opentau.datasets.utils import IterableNamespace
-from opentau.utils.utils import init_logging
-
-
-def run_server(
-    dataset: LeRobotDataset | IterableNamespace | None,
-    episodes: list[int] | None,
-    host: str,
-    port: str,
-    static_folder: Path,
-    template_folder: Path,
-):
-    app = Flask(__name__, static_folder=static_folder.resolve(), template_folder=template_folder.resolve())
-    app.config["SEND_FILE_MAX_AGE_DEFAULT"] = 0  # specifying not to cache
-
-    @app.route("/")
-    def hommepage(dataset=dataset):
-        if dataset:
-            dataset_namespace, dataset_name = dataset.repo_id.split("/")
-            return redirect(
-                url_for(
-                    "show_episode",
-                    dataset_namespace=dataset_namespace,
-                    dataset_name=dataset_name,
-                    episode_id=0,
-                )
-            )
-
-        dataset_param, episode_param = None, None
-        all_params = request.args
-        if "dataset" in all_params:
-            dataset_param = all_params["dataset"]
-        if "episode" in all_params:
-            episode_param = int(all_params["episode"])
-
-        if dataset_param:
-            dataset_namespace, dataset_name = dataset_param.split("/")
-            return redirect(
-                url_for(
-                    "show_episode",
-                    dataset_namespace=dataset_namespace,
-                    dataset_name=dataset_name,
-                    episode_id=episode_param if episode_param is not None else 0,
-                )
-            )
-
-        featured_datasets = [
-            "lerobot/aloha_static_cups_open",
-            "lerobot/columbia_cairlab_pusht_real",
-            "lerobot/taco_play",
-        ]
-        return render_template(
-            "visualize_dataset_homepage.html",
-            featured_datasets=featured_datasets,
-            lerobot_datasets=available_datasets,
-        )
-
-    @app.route("/<string:dataset_namespace>/<string:dataset_name>")
-    def show_first_episode(dataset_namespace, dataset_name):
-        first_episode_id = 0
-        return redirect(
-            url_for(
-                "show_episode",
-                dataset_namespace=dataset_namespace,
-                dataset_name=dataset_name,
-                episode_id=first_episode_id,
-            )
-        )
-
-    @app.route("/<string:dataset_namespace>/<string:dataset_name>/episode_<int:episode_id>")
-    def show_episode(dataset_namespace, dataset_name, episode_id, dataset=dataset, episodes=episodes):
-        repo_id = f"{dataset_namespace}/{dataset_name}"
-        try:
-            if dataset is None:
-                dataset = get_dataset_info(repo_id)
-        except FileNotFoundError:
-            return (
-                "Make sure to convert your LeRobotDataset to v2 & above. See how to convert your dataset at https://github.com/huggingface/lerobot/pull/461",
-                400,
-            )
-        dataset_version = (
-            str(dataset.meta._version) if isinstance(dataset, LeRobotDataset) else dataset.codebase_version
-        )
-        match = re.search(r"v(\d+)\.", dataset_version)
-        if match:
-            major_version = int(match.group(1))
-            if major_version < 2:
-                return "Make sure to convert your LeRobotDataset to v2 & above."
-
-        episode_data_csv_str, columns, ignored_columns = get_episode_data(dataset, episode_id)
-        dataset_info = {
-            "repo_id": f"{dataset_namespace}/{dataset_name}",
-            "num_samples": dataset.num_frames
-            if isinstance(dataset, LeRobotDataset)
-            else dataset.total_frames,
-            "num_episodes": dataset.num_episodes
-            if isinstance(dataset, LeRobotDataset)
-            else dataset.total_episodes,
-            "fps": dataset.fps,
-        }
-        if isinstance(dataset, LeRobotDataset):
-            video_paths = [
-                dataset.meta.get_video_file_path(episode_id, key) for key in dataset.meta.video_keys
-            ]
-            videos_info = [
-                {"url": url_for("static", filename=video_path), "filename": video_path.parent.name}
-                for video_path in video_paths
-            ]
-            tasks = dataset.meta.episodes[episode_id]["tasks"]
-        else:
-            video_keys = [key for key, ft in dataset.features.items() if ft["dtype"] == "video"]
-            videos_info = [
-                {
-                    "url": f"https://huggingface.co/datasets/{repo_id}/resolve/main/"
-                    + dataset.video_path.format(
-                        episode_chunk=int(episode_id) // dataset.chunks_size,
-                        video_key=video_key,
-                        episode_index=episode_id,
-                    ),
-                    "filename": video_key,
-                }
-                for video_key in video_keys
-            ]
-
-            response = requests.get(
-                f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/episodes.jsonl", timeout=5
-            )
-            response.raise_for_status()
-            # Split into lines and parse each line as JSON
-            tasks_jsonl = [json.loads(line) for line in response.text.splitlines() if line.strip()]
-
-            filtered_tasks_jsonl = [row for row in tasks_jsonl if row["episode_index"] == episode_id]
-            tasks = filtered_tasks_jsonl[0]["tasks"]
-
-        videos_info[0]["language_instruction"] = tasks
-
-        if episodes is None:
-            episodes = list(
-                range(dataset.num_episodes if isinstance(dataset, LeRobotDataset) else dataset.total_episodes)
-            )
-
-        return render_template(
-            "visualize_dataset_template.html",
-            episode_id=episode_id,
-            episodes=episodes,
-            dataset_info=dataset_info,
-            videos_info=videos_info,
-            episode_data_csv_str=episode_data_csv_str,
-            columns=columns,
-            ignored_columns=ignored_columns,
-        )
-
-    app.run(host=host, port=port)
-
-
-def get_ep_csv_fname(episode_id: int):
-    ep_csv_fname = f"episode_{episode_id}.csv"
-    return ep_csv_fname
-
-
-def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index):
-    """Get a csv str containing timeseries data of an episode (e.g. state and action).
-    This file will be loaded by Dygraph javascript to plot data in real time."""
-    columns = []
-
-    selected_columns = [col for col, ft in dataset.features.items() if ft["dtype"] in ["float32", "int32"]]
-    selected_columns.remove("timestamp")
-
-    ignored_columns = []
-    for column_name in selected_columns:
-        shape = dataset.features[column_name]["shape"]
-        shape_dim = len(shape)
-        if shape_dim > 1:
-            selected_columns.remove(column_name)
-            ignored_columns.append(column_name)
-
-    # init header of csv with state and action names
-    header = ["timestamp"]
-
-    for column_name in selected_columns:
-        dim_state = (
-            dataset.meta.shapes[column_name][0]
-            if isinstance(dataset, LeRobotDataset)
-            else dataset.features[column_name].shape[0]
-        )
-
-        if "names" in dataset.features[column_name] and dataset.features[column_name]["names"]:
-            column_names = dataset.features[column_name]["names"]
-            while not isinstance(column_names, list):
-                column_names = list(column_names.values())[0]
-        else:
-            column_names = [f"{column_name}_{i}" for i in range(dim_state)]
-        columns.append({"key": column_name, "value": column_names})
-
-        header += column_names
-
-    selected_columns.insert(0, "timestamp")
-
-    if isinstance(dataset, LeRobotDataset):
-        from_idx = dataset.episode_data_index["from"][episode_index]
-        to_idx = dataset.episode_data_index["to"][episode_index]
-        data = (
-            dataset.hf_dataset.select(range(from_idx, to_idx))
-            .select_columns(selected_columns)
-            .with_format("pandas")
-        )
-    else:
-        repo_id = dataset.repo_id
-
-        url = f"https://huggingface.co/datasets/{repo_id}/resolve/main/" + dataset.data_path.format(
-            episode_chunk=int(episode_index) // dataset.chunks_size, episode_index=episode_index
-        )
-        df = pd.read_parquet(url)
-        data = df[selected_columns]  # Select specific columns
-
-    rows = np.hstack(
-        (
-            np.expand_dims(data["timestamp"], axis=1),
-            *[np.vstack(data[col]) for col in selected_columns[1:]],
-        )
-    ).tolist()
-
-    # Convert data to CSV string
-    csv_buffer = StringIO()
-    csv_writer = csv.writer(csv_buffer)
-    # Write header
-    csv_writer.writerow(header)
-    # Write data rows
-    csv_writer.writerows(rows)
-    csv_string = csv_buffer.getvalue()
-
-    return csv_string, columns, ignored_columns
-
-
-def get_episode_video_paths(dataset: LeRobotDataset, ep_index: int) -> list[str]:
-    # get first frame of episode (hack to get video_path of the episode)
-    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
-    return [
-        dataset.hf_dataset.select_columns(key)[first_frame_idx][key]["path"]
-        for key in dataset.meta.video_keys
-    ]
-
-
-def get_episode_language_instruction(dataset: LeRobotDataset, ep_index: int) -> list[str]:
-    # check if the dataset has language instructions
-    if "language_instruction" not in dataset.features:
-        return None
-
-    # get first frame index
-    first_frame_idx = dataset.episode_data_index["from"][ep_index].item()
-
-    language_instruction = dataset.hf_dataset[first_frame_idx]["language_instruction"]
-    # TODO (michel-aractingi) hack to get the sentence, some strings in openx are badly stored
-    # with the tf.tensor appearing in the string
-    return language_instruction.removeprefix("tf.Tensor(b'").removesuffix("', shape=(), dtype=string)")
-
-
-def get_dataset_info(repo_id: str) -> IterableNamespace:
-    response = requests.get(
-        f"https://huggingface.co/datasets/{repo_id}/resolve/main/meta/info.json", timeout=5
-    )
-    response.raise_for_status()  # Raises an HTTPError for bad responses
-    dataset_info = response.json()
-    dataset_info["repo_id"] = repo_id
-    return IterableNamespace(dataset_info)
-
-
-def visualize_dataset_html(
-    dataset: LeRobotDataset | None,
-    episodes: list[int] | None = None,
-    output_dir: Path | None = None,
-    serve: bool = True,
-    host: str = "127.0.0.1",
-    port: int = 9090,
-    force_override: bool = False,
-) -> Path | None:
-    init_logging()
-
-    template_dir = Path(__file__).resolve().parent.parent / "templates"
-
-    if output_dir is None:
-        # Create a temporary directory that will be automatically cleaned up
-        output_dir = tempfile.mkdtemp(prefix="lerobot_visualize_dataset_")
-
-    output_dir = Path(output_dir)
-    if output_dir.exists():
-        if force_override:
-            shutil.rmtree(output_dir)
-        else:
-            logging.info(f"Output directory already exists. Loading from it: '{output_dir}'")
-
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    static_dir = output_dir / "static"
-    static_dir.mkdir(parents=True, exist_ok=True)
-
-    if dataset is None:
-        if serve:
-            run_server(
-                dataset=None,
-                episodes=None,
-                host=host,
-                port=port,
-                static_folder=static_dir,
-                template_folder=template_dir,
-            )
-    else:
-        # Create a simlink from the dataset video folder containing mp4 files to the output directory
-        # so that the http server can get access to the mp4 files.
-        if isinstance(dataset, LeRobotDataset):
-            ln_videos_dir = static_dir / "videos"
-            if not ln_videos_dir.exists():
-                ln_videos_dir.symlink_to((dataset.root / "videos").resolve())
-
-        if serve:
-            run_server(dataset, episodes, host, port, static_dir, template_dir)
-
-
-def create_mock_train_config() -> TrainPipelineConfig:
-    """Create a mock TrainPipelineConfig for dataset visualization.
-
-    Returns:
-        TrainPipelineConfig: A mock config with default values.
-    """
-    return TrainPipelineConfig(
-        dataset_mixture=DatasetMixtureConfig(),  # Will be set by the dataset
-        resolution=(224, 224),
-        num_cams=2,
-        max_state_dim=32,
-        max_action_dim=32,
-        action_chunk=50,
-        loss_weighting={"MSE": 1, "CE": 1},
-        num_workers=4,
-        batch_size=8,
-        steps=100_000,
-        log_freq=200,
-        save_checkpoint=True,
-        save_freq=20_000,
-        use_policy_training_preset=True,
-        wandb=WandBConfig(),
-    )
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument(
-        "--repo-id",
-        type=str,
-        default=None,
-        help="Name of hugging face repositery containing a LeRobotDataset dataset (e.g. `lerobot/pusht` for https://huggingface.co/datasets/lerobot/pusht).",
-    )
-    parser.add_argument(
-        "--root",
-        type=Path,
-        default=None,
-        help="Root directory for a dataset stored locally (e.g. `--root data`). By default, the dataset will be loaded from hugging face cache folder, or downloaded from the hub if available.",
-    )
-    parser.add_argument(
-        "--load-from-hf-hub",
-        type=int,
-        default=0,
-        help="Load videos and parquet files from HF Hub rather than local system.",
-    )
-    parser.add_argument(
-        "--episodes",
-        type=int,
-        nargs="*",
-        default=None,
-        help="Episode indices to visualize (e.g. `0 1 5 6` to load episodes of index 0, 1, 5 and 6). By default loads all episodes.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=Path,
-        default=None,
-        help="Directory path to write html files and kickoff a web server. By default write them to 'outputs/visualize_dataset/REPO_ID'.",
-    )
-    parser.add_argument(
-        "--serve",
-        type=int,
-        default=1,
-        help="Launch web server.",
-    )
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="127.0.0.1",
-        help="Web host used by the http server.",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=9090,
-        help="Web port used by the http server.",
-    )
-    parser.add_argument(
-        "--force-override",
-        type=int,
-        default=0,
-        help="Delete the output directory if it exists already.",
-    )
-
-    parser.add_argument(
-        "--tolerance-s",
-        type=float,
-        default=1e-4,
-        help=(
-            "Tolerance in seconds used to ensure data timestamps respect the dataset fps value"
-            "This is argument passed to the constructor of LeRobotDataset and maps to its tolerance_s constructor argument"
-            "If not given, defaults to 1e-4."
-        ),
-    )
-
-    args = parser.parse_args()
-    kwargs = vars(args)
-    repo_id = kwargs.pop("repo_id")
-    load_from_hf_hub = kwargs.pop("load_from_hf_hub")
-    root = kwargs.pop("root")
-    tolerance_s = kwargs.pop("tolerance_s")
-
-    dataset = None
-    if repo_id:
-        dataset = (
-            LeRobotDataset(create_mock_train_config(), repo_id, root=root, tolerance_s=tolerance_s)
-            if not load_from_hf_hub
-            else get_dataset_info(repo_id)
-        )
-
-    visualize_dataset_html(dataset, **vars(args))
-
-
-if __name__ == "__main__":
-    main()