diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index d07a2767d..a7c0b774e 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -139,9 +139,27 @@ def visualize_dataset( logging.info("Logging to Rerun") first_index = None + series_names_logged = False for batch in tqdm.tqdm(dataloader, total=len(dataloader)): if first_index is None: first_index = batch["index"][0].item() + + # Name each series once (static) so all dimensions share a single view while keeping labels. + if not series_names_logged: + if ACTION in batch: + rr.log( + ACTION, + rr.SeriesLines(names=[f"{ACTION}_{d}" for d in range(batch[ACTION].shape[-1])]), + static=True, + ) + if OBS_STATE in batch: + rr.log( + "state", + rr.SeriesLines(names=[f"state_{d}" for d in range(batch[OBS_STATE].shape[-1])]), + static=True, + ) + series_names_logged = True + # iterate over the batch for i in range(len(batch["index"])): rr.set_time("frame_index", sequence=batch["index"][i].item() - first_index) @@ -155,13 +173,11 @@ 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.Scalars(val.item())) + rr.log(ACTION, rr.Scalars(batch[ACTION][i].numpy())) # display each dimension of observed state space (e.g. agent position in joint space) if OBS_STATE in batch: - for dim_idx, val in enumerate(batch[OBS_STATE][i]): - rr.log(f"state/{dim_idx}", rr.Scalars(val.item())) + rr.log("state", rr.Scalars(batch[OBS_STATE][i].numpy())) if DONE in batch: rr.log(DONE, rr.Scalars(batch[DONE][i].item())) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index d9d5bf6b5..117316450 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -76,8 +76,9 @@ def log_rerun_data( - Scalars values (floats, ints) are logged as `rr.Scalars`. - 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`. - - 1D NumPy arrays are logged as a series of individual scalars, with each element indexed. - - Other multi-dimensional arrays are flattened and logged as individual scalars. + - 1D NumPy arrays are logged as a single `rr.Scalars` batch under one entity path, so that every + dimension shares the same view instead of being split across one view per element. + - Other multi-dimensional arrays are flattened and logged as a single `rr.Scalars` batch. Keys are automatically namespaced with "observation." or "action." if not already present. @@ -90,6 +91,15 @@ def log_rerun_data( require_package("rerun-sdk", extra="viz", import_name="rerun") import rerun as rr + def _log_vector(key: str, arr: np.ndarray) -> None: + """ + Logs an array as one batch so all dimensions share a single view, while keeping the per-dimension `{key}_{i}` names via SeriesLines. + Arrays with more than one dimension are flattened. + """ + arr = arr.reshape(-1).astype(float) + rr.log(key, rr.SeriesLines(names=[f"{key}_{i}" for i in range(arr.shape[0])]), static=True) + rr.log(key, rr.Scalars(arr)) + if observation: for k, v in observation.items(): if v is None: @@ -104,8 +114,7 @@ def log_rerun_data( if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): arr = np.transpose(arr, (1, 2, 0)) if arr.ndim == 1: - for i, vi in enumerate(arr): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + _log_vector(key, arr) else: img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True) @@ -119,11 +128,4 @@ def log_rerun_data( if _is_scalar(v): rr.log(key, rr.Scalars(float(v))) elif isinstance(v, np.ndarray): - if v.ndim == 1: - for i, vi in enumerate(v): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) - else: - # Fall back to flattening higher-dimensional arrays - flat = v.flatten() - for i, vi in enumerate(flat): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + _log_vector(key, v)