mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-14 08:09:45 +00:00
fix(rerun audio): fixing rerun visualization for audio
This commit is contained in:
@@ -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))
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)))
|
||||
|
||||
Reference in New Issue
Block a user