From edb5559b5b9116e752d0d14a7759e66ff063a9da Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Fri, 31 Oct 2025 18:39:19 +0100 Subject: [PATCH] fix(rerun audio): fixing rerun visualization for audio --- examples/lekiwi/teleoperate.py | 3 +- examples/phone_to_so100/teleoperate.py | 5 +-- examples/so100_to_so100_EE/teleoperate.py | 7 ++-- src/lerobot/scripts/lerobot_teleoperate.py | 1 + src/lerobot/utils/visualization_utils.py | 41 ++++++++++++---------- 5 files changed, 33 insertions(+), 24 deletions(-) diff --git a/examples/lekiwi/teleoperate.py b/examples/lekiwi/teleoperate.py index 102bd31f5..1983abcd6 100644 --- a/examples/lekiwi/teleoperate.py +++ b/examples/lekiwi/teleoperate.py @@ -49,6 +49,7 @@ def main(): raise ValueError("Robot or teleop is not connected!") print("Starting teleop loop...") + start = time.perf_counter() while True: t0 = time.perf_counter() @@ -69,7 +70,7 @@ def main(): _ = robot.send_action(action) # Visualize - log_rerun_data(observation=observation, action=action) + log_rerun_data(observation=observation, action=action, log_time=time.perf_counter() - start) precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) diff --git a/examples/phone_to_so100/teleoperate.py b/examples/phone_to_so100/teleoperate.py index 6eaaec806..227657cda 100644 --- a/examples/phone_to_so100/teleoperate.py +++ b/examples/phone_to_so100/teleoperate.py @@ -89,12 +89,13 @@ def main(): teleop_device.connect() # Init rerun viewer - init_rerun(session_name="phone_so100_teleop") + init_rerun(session_name="phone_so100_teleop", robot=robot, reset_time=True) if not robot.is_connected or not teleop_device.is_connected: raise ValueError("Robot or teleop is not connected!") print("Starting teleop loop. Move your phone to teleoperate the robot...") + start = time.perf_counter() while True: t0 = time.perf_counter() @@ -111,7 +112,7 @@ def main(): _ = robot.send_action(joint_action) # Visualize - log_rerun_data(observation=phone_obs, action=joint_action) + log_rerun_data(observation=phone_obs, action=joint_action, log_time=time.perf_counter() - start) precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) diff --git a/examples/so100_to_so100_EE/teleoperate.py b/examples/so100_to_so100_EE/teleoperate.py index 71d2899de..a071557d3 100644 --- a/examples/so100_to_so100_EE/teleoperate.py +++ b/examples/so100_to_so100_EE/teleoperate.py @@ -94,9 +94,10 @@ def main(): leader.connect() # Init rerun viewer - init_rerun(session_name="so100_so100_EE_teleop") + init_rerun(session_name="so100_so100_EE_teleop", robot=follower, reset_time=True) print("Starting teleop loop...") + start = time.perf_counter() while True: t0 = time.perf_counter() @@ -116,7 +117,9 @@ def main(): _ = follower.send_action(follower_joints_act) # Visualize - log_rerun_data(observation=leader_ee_act, action=follower_joints_act) + log_rerun_data( + observation=leader_ee_act, action=follower_joints_act, log_time=time.perf_counter() - start + ) precise_sleep(max(1.0 / FPS - (time.perf_counter() - t0), 0.0)) diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index 80bebacf3..ae8a7c3ab 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -187,6 +187,7 @@ def teleop_loop( observation=obs_transition, action=teleop_action, compress_images=display_compressed_images, + log_time=time.perf_counter() - start, ) print("\n" + "-" * (display_len + 10)) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index fd83391f3..8472203cc 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/src/lerobot/utils/visualization_utils.py @@ -14,6 +14,7 @@ import numbers import os +import time from uuid import uuid4 import numpy as np @@ -48,8 +49,9 @@ def init_rerun( rr.init( application_id=session_name, recording_id=uuid4(), - default_blueprint=build_rerun_blueprint(robot) if robot is not None else None, ) + if robot is not None: + rr.send_blueprint(build_rerun_blueprint(robot)) memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") if ip and port: rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") @@ -57,7 +59,7 @@ def init_rerun( rr.spawn(memory_limit=memory_limit) if reset_time: - rr.set_time_seconds("episode_time", seconds=0.0) + rr.set_time("episode_time", timestamp=0.0) def _is_scalar(x): @@ -80,26 +82,26 @@ def build_rerun_blueprint(robot: Robot) -> rr.blueprint.Grid: """ contents = [ rr.blueprint.TimeSeriesView( - origin="states_actions", + origin="data", plot_legend=rr.blueprint.PlotLegend(visible=True), ) ] if robot.microphones: contents += [ rr.blueprint.TimeSeriesView( - origin="microphones", + origin="audio", plot_legend=rr.blueprint.PlotLegend(visible=True), ) ] if robot.cameras: contents += [ rr.blueprint.Spatial2DView( - origin=camera_name, + origin=OBS_PREFIX + camera_name, ) for camera_name in robot.cameras ] - return rr.blueprint.Grid(contents) + return rr.blueprint.Grid(*contents) def log_rerun_data( @@ -128,8 +130,9 @@ def log_rerun_data( log_time: The time to log the data in the "episode_time" timeline. If None, the current time is used in Rerun's default timeline. """ - if log_time is not None: - rr.set_time_seconds("episode_time", seconds=log_time) + if log_time is None: + log_time = time.perf_counter() + rr.set_time("episode_time", timestamp=log_time) if observation: for k, v in observation.items(): @@ -138,29 +141,29 @@ def log_rerun_data( key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}" if _is_scalar(v): - rr.log(key, rr.Scalars(float(v))) + rr.log("data/" + key, rr.Scalars(float(v))) elif isinstance(v, np.ndarray): arr = v # Convert CHW -> HWC when needed 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)) - # Convert channel x samples -> samples x channel when needed - elif arr.ndim == 2 and arr.shape[0] < arr.shape[1]: + # Convert samples x channels -> channels x samples when needed + elif arr.ndim == 2 and arr.shape[1] < arr.shape[0]: arr = np.transpose(arr, (1, 0)) if arr.ndim == 1: for i, vi in enumerate(arr): - rr.log(f"{key}_{i}", rr.Scalars(float(vi))) + rr.log("data/" + f"{key}_{i}", rr.Scalars(float(vi))) elif arr.ndim == 2: - for i, channel_arr in enumerate(arr.T): + for i, channel_arr in enumerate(arr): rr.send_columns( "audio/" + key + f"_channel_{i}", # TODO(CarolinePascal): Get actual channel number/name indexes=[ - rr.TimeSecondsColumn( + rr.TimeColumn( "episode_time", - times=log_time + timestamp=log_time + np.linspace( -DEFAULT_AUDIO_CHUNK_DURATION, 0, @@ -169,7 +172,7 @@ def log_rerun_data( ), ) ], - columns=rr.Scalar.columns(scalar=channel_arr), + columns=rr.Scalars.columns(scalars=channel_arr), ) elif arr.ndim == 3: rr.log(key, rr.Image(arr), static=True) @@ -184,13 +187,13 @@ def log_rerun_data( key = k if str(k).startswith(ACTION_PREFIX) else f"{ACTION}.{k}" if _is_scalar(v): - rr.log(key, rr.Scalars(float(v))) + rr.log("data/" + 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))) + rr.log("data/" + 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))) + rr.log("data/" + f"{key}_{i}", rr.Scalars(float(vi)))