diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index 66ac5a511..a168d716b 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -113,36 +113,26 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: def build_blueprint_from_dataset(dataset: LeRobotDataset): """Build a Rerun blueprint laying out camera images and time series for the given dataset. - Camera images are arranged in a grid on the left, and the available scalar signals - (action, state, reward, done, success) are stacked as time series views on the right. - The per-dimension series names and colors for ``action`` and ``state`` are applied - directly via blueprint overrides. + Camera images and scalar signals (action, state, reward, done, success) are arranged in a grid. + The per-dimension series names and colors for ``action`` and ``state`` are applied directly + via blueprint overrides. """ import rerun as rr import rerun.blueprint as rrb - image_views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys] + views = [rrb.Spatial2DView(origin=key, name=key) for key in dataset.meta.camera_keys] - timeseries_views = [] # Style multi-dimensional signals (action, state) with per-dimension names and colors. for origin, key in ((ACTION, ACTION), ("state", OBS_STATE)): if key in dataset.features: names = get_feature_names(dataset, key) styling = rr.SeriesLines(names=names, colors=get_sequential_colors(len(names))) - timeseries_views.append( - rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling}) - ) + views.append(rrb.TimeSeriesView(origin=origin, name=origin, overrides={origin: styling})) for key in (DONE, REWARD, "next.success"): if key in dataset.features: - timeseries_views.append(rrb.TimeSeriesView(origin=key, name=key)) + views.append(rrb.TimeSeriesView(origin=key, name=key)) - contents = [] - if image_views: - contents.append(rrb.Grid(*image_views, name="images")) - if timeseries_views: - contents.append(rrb.Vertical(*timeseries_views, name="time series")) - - return rrb.Blueprint(rrb.Horizontal(*contents) if contents else rrb.Grid()) + return rrb.Blueprint(rrb.Grid(*views)) def visualize_dataset( diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index a9da5f940..1ff99a482 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -66,25 +66,18 @@ def _is_scalar(x): def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]): """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views. - Camera images are arranged in a grid on the left, and the observation and action scalars are stacked as time series views on the right. + Camera images, observation and action scalars are arranged in a grid. """ import rerun.blueprint as rrb - image_views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)] + views = [rrb.Spatial2DView(origin=path, name=path) for path in sorted(image_paths)] - timeseries_views = [] if observation_paths: - timeseries_views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths))) + views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths))) if action_paths: - timeseries_views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths))) + views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths))) - contents = [] - if image_views: - contents.append(rrb.Grid(*image_views, name="images")) - if timeseries_views: - contents.append(rrb.Vertical(*timeseries_views, name="time series")) - - return rrb.Blueprint(rrb.Horizontal(*contents) if contents else rrb.Grid()) + return rrb.Blueprint(rrb.Grid(*views)) def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str], rr) -> None: