opentau 0.1.0__py3-none-any.whl

opentau/scripts/train.py

@@ -0,0 +1,379 @@
+#!/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.
+import json
+import logging
+import os
+from contextlib import nullcontext
+from pprint import pformat
+from typing import Any
+
+import accelerate
+import torch
+from accelerate.optimizer import AcceleratedOptimizer
+from accelerate.scheduler import AcceleratedScheduler
+from accelerate.utils import DistributedDataParallelKwargs, gather_object
+from termcolor import colored
+
+from opentau.configs import parser
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.factory import make_dataset_mixture
+from opentau.datasets.utils import cycle
+from opentau.envs.factory import make_envs
+from opentau.envs.utils import close_envs
+from opentau.optim.factory import make_optimizer_and_scheduler
+from opentau.policies.factory import make_policy
+from opentau.policies.pretrained import PreTrainedPolicy
+from opentau.scripts.eval import consolidate_eval_info, eval_policy_all
+from opentau.utils.accelerate_utils import set_proc_accelerator
+from opentau.utils.logging_utils import AverageMeter, MetricsTracker
+from opentau.utils.random_utils import set_seed
+from opentau.utils.train_utils import (
+    get_step_checkpoint_dir,
+    get_step_identifier,
+    load_training_state,
+    load_training_step,
+    prune_old_checkpoints,
+    save_checkpoint,
+)
+from opentau.utils.utils import (
+    encode_accelerator_state_dict,
+    format_big_number,
+    init_logging,
+    is_launched_with_accelerate,
+)
+
+
+def update_policy(
+    train_config: TrainPipelineConfig,
+    train_metrics: MetricsTracker,
+    policy: PreTrainedPolicy,
+    batch: Any,
+    optimizer: AcceleratedOptimizer,
+    grad_clip_norm: float,
+    accelerator: accelerate.Accelerator,
+    lr_scheduler: AcceleratedScheduler | None = None,
+) -> tuple[MetricsTracker, dict]:
+    policy.train()
+    losses = policy.forward(batch)
+    loss = (
+        train_config.loss_weighting["MSE"] * losses["MSE"] + train_config.loss_weighting["CE"] * losses["CE"]
+    )
+
+    # accelerator.backward(loss)
+    # accelerator.unscale_gradients(optimizer=optimizer)
+
+    # if accelerator.sync_gradients:
+    #     grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
+    #     if accelerator.is_main_process:
+    #         train_metrics.grad_norm = grad_norm
+
+    # optimizer.step()
+    # optimizer.zero_grad()
+
+    # Step through pytorch scheduler at every batch instead of epoch
+    if lr_scheduler is not None:
+        lr_scheduler.step()
+
+    # This calls `torch.distributed.all_gather_into_tensor` under the hood, which is not so efficient.
+    # We don't actually want to broadcast the gathered tensors to all processes, but only to the main process.
+    # Nonetheless, we still do this for correctness, safety, and simplicity.
+    _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()
+    )
+    # This actually calls `.update` method of the `AverageMeter` class. This operation is not idempotent.
+    # See MetricsTracker.__setattr__ for more details.
+    # In other words, setting `train_metrics.loss = 1` and `train_metrics.loss = 2` consecutively results in
+    #   an average of 1.5 when formatted as a string, not just 2.
+    if accelerator.is_main_process:
+        train_metrics.loss = loss
+        train_metrics.mse_loss = mse_loss
+        train_metrics.ce_loss = ce_loss
+        train_metrics.l1_loss = l1_loss
+        train_metrics.accuracy = accuracy
+        train_metrics.lr = optimizer.param_groups[0]["lr"]
+
+    return train_metrics
+
+
+@parser.wrap()
+def train(cfg: TrainPipelineConfig):
+    cfg.validate()
+
+    accelerator_kwargs = {
+        "step_scheduler_with_optimizer": False,
+        "split_batches": False,  # split_batches == True is not working anyways
+        "kwargs_handlers": [DistributedDataParallelKwargs(find_unused_parameters=True)],
+    }
+    if cfg.wandb.enable:
+        accelerator_kwargs["log_with"] = "wandb"
+    if cfg.gradient_accumulation_steps > 1:
+        accelerator_kwargs["gradient_accumulation_steps"] = cfg.gradient_accumulation_steps
+
+    accelerator = accelerate.Accelerator(**accelerator_kwargs)
+    init_logging(accelerator, level=logging.DEBUG if cfg.debug else logging.INFO)
+    # Register accelerator globally for use in other modules, (e.g., detect current rank, etc.)
+    set_proc_accelerator(accelerator)
+
+    logging.info(pformat(cfg.to_dict()))
+
+    if accelerator.is_main_process:
+        accelerator_config = encode_accelerator_state_dict(accelerator.state.__dict__)
+        logging.info(pformat(accelerator_config))
+
+        # Ensure `gradient_accumulation_steps` is consistent between TrainPipelineConfig and DeepSpeedConfig
+        if accelerator.distributed_type == accelerate.DistributedType.DEEPSPEED:
+            deepspeed_config, deepspeed_key = accelerator.deepspeed_plugin.hf_ds_config.find_config_node(
+                "gradient_accumulation_steps"
+            )
+            ds_grad_acc_steps = deepspeed_config.get(deepspeed_key, 1)
+            if ds_grad_acc_steps != cfg.gradient_accumulation_steps:
+                raise ValueError(
+                    "The `gradient_accumulation_steps` in TrainPipelineConfig does not match the value "
+                    f"specified in DeepSpeedConfig {cfg.gradient_accumulation_steps} != {ds_grad_acc_steps}. "  # nosec B608
+                )
+
+        if cfg.wandb.enable:
+            step = load_training_step(cfg.checkpoint_path) if cfg.resume else None
+            slurm_dict = {k: v for k, v in os.environ.items() if k.startswith("SLURM_")}
+            accelerator.init_trackers(
+                cfg.wandb.project,
+                config={**cfg.to_dict(), "accelerator": accelerator_config, "slurm": slurm_dict},
+                init_kwargs={"wandb": cfg.wandb.to_wandb_kwargs(step=step)},
+            )
+            tracker = accelerator.get_tracker("wandb", unwrap=True)
+            cfg.wandb.run_id = tracker.id
+            logging.info(f"tracker initialized with wandb job id: {tracker.id}")
+
+    if cfg.seed is not None:
+        set_seed(cfg.seed, accelerator=accelerator)
+
+    # Enable anomaly detection for debugging NaN/Inf values
+    # (warning: large computational overhead)
+    torch.autograd.set_detect_anomaly(cfg.trace_nans)
+    if cfg.trace_nans:
+        logging.warning("Anomaly detection is enabled. This may significantly slow down training.")
+    else:
+        logging.info("Anomaly detection is disabled.")
+
+    logging.info("Creating dataset")
+    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,
+    eval_envs = None
+    if cfg.eval_freq > 0 and cfg.env is not None:
+        logging.info("Creating env")
+        eval_envs = make_envs(
+            cfg.env, cfg, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs
+        )
+
+    logging.info("Creating policy")
+    policy = make_policy(cfg=cfg.policy, ds_meta=dataset.meta)
+    policy.to(torch.bfloat16)
+    logging.info("Creating optimizer and scheduler")
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
+
+    step = 0  # number of policy updates (forward + backward + optim)
+
+    if accelerator.is_main_process:
+        num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad)
+        num_total_params = sum(p.numel() for p in policy.parameters())
+        logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
+        logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})")
+        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)
+
+    # Register the LR scheduler for checkpointing
+    accelerator.register_for_checkpointing(lr_scheduler)
+
+    if cfg.resume:
+        # load accelerator state
+        # This will load the model, optimizer, and lr_scheduler state
+        accelerator.load_state(cfg.checkpoint_path)
+
+        # all processes should load the step & rng states
+        step = load_training_state(cfg.checkpoint_path)
+        logging.info(f"Resuming training from checkpoint {cfg.checkpoint_path}")
+
+    policy.train()
+
+    # setup metrics tracker to average metrics over the logging interval
+    train_metrics = {
+        "loss": AverageMeter("total_loss", ":.3f"),
+        "mse_loss": AverageMeter("mse_loss", ":.3f"),
+        "ce_loss": AverageMeter("ce_loss", ":.3f"),
+        "l1_loss": AverageMeter("l1_loss", ":.3f"),
+        "accuracy": AverageMeter("accuracy", ":.3f"),
+        "lr": AverageMeter("lr", ":0.1e"),
+        "grad_norm": AverageMeter("grad_norm", ":.3f"),
+    }
+    train_tracker = MetricsTracker(
+        cfg.batch_size * accelerator.num_processes,  # split_batches are not working
+        train_metrics,
+        initial_step=step,
+    )
+
+    if accelerator.is_main_process:
+        logging.info("Start offline training on a fixed dataset")
+
+    for _ in range(step, cfg.steps):
+        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)
+
+                train_tracker = update_policy(
+                    cfg,
+                    train_tracker,
+                    policy,
+                    batch,
+                    optimizer,
+                    cfg.optimizer.grad_clip_norm,
+                    accelerator=accelerator,
+                    lr_scheduler=lr_scheduler,
+                )
+
+        # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
+        # increment `step` here.
+        step += 1
+        train_tracker.step()
+        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
+
+        # 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)
+            train_tracker.reset_averages()
+
+        if is_saving_step:
+            # TODO: investigate whether this barrier is needed
+            accelerator.wait_for_everyone()
+            checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
+
+            # save the accelerator state
+            # This will save the model, optimizer, and lr_scheduler state
+            accelerator.save_state(checkpoint_dir)
+
+            # save axillary objects such as configs, training step, and rng state
+            if accelerator.is_main_process:
+                logging.info(f"Checkpoint policy after step {step}")
+                cfg.policy.pretrained_path = checkpoint_dir
+                save_checkpoint(checkpoint_dir, step, cfg)
+                if cfg.last_checkpoint_only:
+                    prune_old_checkpoints(checkpoint_dir)
+
+            # This barrier is probably necessary to ensure
+            # other processes wait for the main process to finish saving
+            accelerator.wait_for_everyone()
+
+        if is_eval_step and eval_envs:
+            step_id = get_step_identifier(step, cfg.steps)
+            logging.info(f"Eval policy at step {step}")
+            with (
+                torch.no_grad(),
+                torch.autocast(device_type=accelerator.device.type) if cfg.policy.use_amp else nullcontext(),
+            ):
+                eval_info = eval_policy_all(
+                    eval_envs,
+                    policy,
+                    cfg.eval.n_episodes,
+                    cfg,
+                    videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}",
+                    max_episodes_rendered=cfg.eval.max_episodes_rendered,
+                    start_seed=cfg.seed,
+                    max_parallel_tasks=cfg.env.max_parallel_tasks,
+                )
+
+            eval_info = gather_object([eval_info])  # gather across all accelerator processes
+            if accelerator.is_main_process:
+                eval_info = consolidate_eval_info(eval_info)
+                # overall metrics (suite-agnostic)
+                aggregated = eval_info["overall"]
+
+                # optional: per-suite logging
+                for suite, suite_info in eval_info.items():
+                    logging.info("Suite %s aggregated: %s", suite, suite_info)
+
+                # meters/tracker
+                eval_metrics = {
+                    "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"),
+                    "pc_success": AverageMeter("success", ":.1f"),
+                    "eval_per_gpu_s": AverageMeter("eval_per_gpu_s", ":.3f"),
+                }
+                eval_tracker = MetricsTracker(
+                    cfg.batch_size,
+                    eval_metrics,
+                    initial_step=step,
+                )
+                eval_tracker.eval_per_gpu_s = aggregated.get("eval_per_gpu_s", float("nan"))
+                eval_tracker.avg_sum_reward = aggregated.get("avg_sum_reward", float("nan"))
+                eval_tracker.pc_success = aggregated.get("pc_success", float("nan"))
+                logging.info(eval_tracker)
+                eval_dict = eval_tracker.to_dict(use_avg=True)
+                accelerator.log({"Success Rate": eval_dict["pc_success"]}, step=step)
+                accelerator.log({"Evaluation Time": eval_dict["eval_per_gpu_s"]}, step=step)
+                for group, v in eval_info["per_group"].items():
+                    accelerator.log({f"Success/{group}": v["pc_success"]}, step=step)
+
+                # Save eval_info to the same directory as videos
+                videos_dir = cfg.output_dir / "eval" / f"videos_step_{step_id}"
+                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()
+
+    if cfg.eval_freq > 0 and eval_envs:
+        close_envs(eval_envs)
+
+    accelerator.end_training()
+    if accelerator.is_main_process:
+        logging.info("End of training")
+
+
+if __name__ == "__main__":
+    if not is_launched_with_accelerate():
+        raise Exception(
+            "This script should be launched with accelerate. Please use `accelerate launch` to run this script."
+        )
+
+    train()

opentau/scripts/visualize_dataset.py

@@ -0,0 +1,294 @@
+#!/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 doesn't 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
+lossy 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.
+
+Examples:
+
+- Visualize data stored on a local machine:
+```
+local$ python src/opentau/scripts/visualize_dataset.py \
+    --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
+
+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
+```
+
+"""
+
+import argparse
+import gc
+import logging
+import time
+from pathlib import Path
+from typing import Iterator
+
+import numpy as np
+import rerun as rr
+import torch
+import torch.utils.data
+import tqdm
+
+from opentau.datasets.lerobot_dataset import LeRobotDataset
+from opentau.scripts.visualize_dataset_html import create_mock_train_config
+
+
+class EpisodeSampler(torch.utils.data.Sampler):
+    def __init__(self, dataset: LeRobotDataset, episode_index: int):
+        from_idx = dataset.episode_data_index["from"][episode_index].item()
+        to_idx = dataset.episode_data_index["to"][episode_index].item()
+        self.frame_ids = range(from_idx, to_idx)
+
+    def __iter__(self) -> Iterator:
+        return iter(self.frame_ids)
+
+    def __len__(self) -> int:
+        return len(self.frame_ids)
+
+
+def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray:
+    assert chw_float32_torch.dtype == torch.float32
+    assert chw_float32_torch.ndim == 3
+    c, h, w = chw_float32_torch.shape
+    assert c < h and c < w, f"expect channel first images, but instead {chw_float32_torch.shape}"
+    hwc_uint8_numpy = (chw_float32_torch * 255).type(torch.uint8).permute(1, 2, 0).numpy()
+    return hwc_uint8_numpy
+
+
+def visualize_dataset(
+    dataset: LeRobotDataset,
+    episode_index: int,
+    batch_size: int = 32,
+    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:
+    if save:
+        assert output_dir is not None, (
+            "Set an output directory where to write .rrd files with `--output-dir path/to/directory`."
+        )
+
+    repo_id = dataset.repo_id
+
+    logging.info("Loading dataloader")
+    episode_sampler = EpisodeSampler(dataset, episode_index)
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        num_workers=num_workers,
+        batch_size=batch_size,
+        sampler=episode_sampler,
+    )
+
+    logging.info("Starting Rerun")
+
+    if mode not in ["local", "distant"]:
+        raise ValueError(mode)
+
+    spawn_local_viewer = mode == "local" and not save
+    rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer)
+
+    # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush
+    # when iterating on a dataloader with `num_workers` > 0
+    # TODO(rcadene): remove `gc.collect` when rerun version 0.16 is out, which includes a fix
+    gc.collect()
+
+    if mode == "distant":
+        rr.serve(open_browser=False, web_port=web_port, ws_port=ws_port)
+
+    logging.info("Logging to Rerun")
+
+    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())
+
+            # display each camera image
+            for key in dataset.meta.camera_keys:
+                # TODO(rcadene): add `.compress()`? is it lossless?
+                rr.log(key, rr.Image(to_hwc_uint8_numpy(batch[key][i])))
+
+            # 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()))
+
+            # 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()))
+
+            if "next.done" in batch:
+                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()))
+
+            if "next.success" in batch:
+                rr.log("next.success", rr.Scalar(batch["next.success"][i].item()))
+
+    if mode == "local" and save:
+        # save .rrd locally
+        output_dir = Path(output_dir)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        repo_id_str = repo_id.replace("/", "_")
+        rrd_path = output_dir / f"{repo_id_str}_episode_{episode_index}.rrd"
+        rr.save(rrd_path)
+        return rrd_path
+
+    elif mode == "distant":
+        # stop the process from exiting since it is serving the websocket connection
+        try:
+            while True:
+                time.sleep(1)
+        except KeyboardInterrupt:
+            print("Ctrl-C received. Exiting.")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Name of hugging face repository containing a LeRobotDataset dataset (e.g. `lerobot/pusht`).",
+    )
+    parser.add_argument(
+        "--episode-index",
+        type=int,
+        required=True,
+        help="Episode to visualize.",
+    )
+    parser.add_argument(
+        "--root",
+        type=Path,
+        default=None,
+        help="Root directory for the 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(
+        "--output-dir",
+        type=Path,
+        default=None,
+        help="Directory path to write a .rrd file when `--save 1` is set.",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=32,
+        help="Batch size loaded by DataLoader.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of processes of Dataloader for loading the data.",
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        default="local",
+        help=(
+            "Mode of viewing between 'local' or 'distant'. "
+            "'local' requires data to be on a local machine. It spawns a viewer to visualize the data locally. "
+            "'distant' creates a server on the distant machine where the data is stored. "
+            "Visualize the data by connecting to the server with `rerun ws://localhost:PORT` on the local machine."
+        ),
+    )
+    parser.add_argument(
+        "--web-port",
+        type=int,
+        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,
+        default=0,
+        help=(
+            "Save a .rrd file in the directory provided by `--output-dir`. "
+            "It also deactivates the spawning of a viewer. "
+            "Visualize the data by running `rerun path/to/file.rrd` on your local machine."
+        ),
+    )
+
+    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")
+    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)
+
+    visualize_dataset(dataset, **vars(args))
+
+
+if __name__ == "__main__":
+    main()