diff --git a/lerobot/scripts/visualize_dataset_html.py b/lerobot/scripts/visualize_dataset_html.py index ac91f0c84..a0da08690 100644 --- a/lerobot/scripts/visualize_dataset_html.py +++ b/lerobot/scripts/visualize_dataset_html.py @@ -245,16 +245,17 @@ def get_episode_data(dataset: LeRobotDataset | IterableNamespace, episode_index) if isinstance(dataset, LeRobotDataset) else dataset.features[column_name].shape[0] ) - header += [f"{column_name}_{i}" for i in range(dim_state)] 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"motor_{i}" for i in range(dim_state)] + 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): diff --git a/lerobot/templates/visualize_dataset_template.html b/lerobot/templates/visualize_dataset_template.html index e5a2f82c6..d81ce630e 100644 --- a/lerobot/templates/visualize_dataset_template.html +++ b/lerobot/templates/visualize_dataset_template.html @@ -246,14 +246,16 @@