opentau 0.1.1__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

@@ -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
 ```
 
 """
@@ -65,7 +52,9 @@ local$ rerun ws://localhost:9087
 import argparse
 import gc
 import logging
+import os
 import time
+import warnings
 from pathlib import Path
 from typing import Iterator
 
@@ -75,8 +64,80 @@ 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
+
+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.
+
+    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,9 +169,9 @@ 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,
+    urdf: Path | None = None,
 ) -> Path | None:
     if save:
         assert output_dir is not None, (
@@ -141,16 +202,27 @@ 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(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")
 
     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:
@@ -160,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
@@ -194,7 +272,7 @@ def visualize_dataset(
             print("Ctrl-C received. Exiting.")
 
 
-def main():
+def parse_args() -> dict:
     parser = argparse.ArgumentParser()
 
     parser.add_argument(
@@ -250,12 +328,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,
@@ -266,7 +338,6 @@ def main():
             "Visualize the data by running `rerun path/to/file.rrd` on your local machine."
         ),
     )
-
     parser.add_argument(
         "--tolerance-s",
         type=float,
@@ -277,17 +348,49 @@ def main():
             "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()
-    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")
+    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(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__":