fix(rerun audio): fixing rerun visualization for audio

This commit is contained in:
CarolinePascal
2025-10-31 18:39:19 +01:00
parent 552ec76195
commit edb5559b5b
5 changed files with 33 additions and 24 deletions
+2 -1
View File
@@ -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))
+3 -2
View File
@@ -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))
+5 -2
View File
@@ -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))
@@ -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))
+22 -19
View File
@@ -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)))