admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-26 22:20:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(visualization): allow remote viewer + compress rerun images (#2756)

Browse Source 
* feat(visualization): allow remote viewer + compress rerun images

* fix(tests): allow named argument in mocked rerun

* feat(visualization): ip instead or url & cli arg for compressing images
This commit is contained in:
Steven Palma
2026-01-07 17:30:45 +01:00
committed by GitHub
parent 7e9d05a799
commit f844c7a458
5 changed files with 69 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/utils/visualization_utils.py
+20 -5
View File
@@ -22,13 +22,25 @@ import rerun as rr
from .constants import OBS_PREFIX, OBS_STR
def init_rerun(session_name: str = "lerobot_control_loop") -> None:
"""Initializes the Rerun SDK for visualizing the control loop."""
def init_rerun(
session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
) -> None:
"""
Initializes the Rerun SDK for visualizing the control loop.
Args:
session_name: Name of the Rerun session.
ip: Optional IP for connecting to a Rerun server.
port: Optional port for connecting to a Rerun server.
"""
batch_size = os.getenv("RERUN_FLUSH_NUM_BYTES", "8000")
os.environ["RERUN_FLUSH_NUM_BYTES"] = batch_size
rr.init(session_name)
memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
rr.spawn(memory_limit=memory_limit)
if ip and port:
rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
else:
rr.spawn(memory_limit=memory_limit)
def _is_scalar(x):
@@ -40,6 +52,7 @@ def _is_scalar(x):
def log_rerun_data(
observation: dict[str, Any] | None = None,
action: dict[str, Any] | None = None,
compress_images: bool = False,
) -> None:
"""
Logs observation and action data to Rerun for real-time visualization.
@@ -48,7 +61,7 @@ def log_rerun_data(
to the Rerun viewer. It handles different data types appropriately:
- Scalars values (floats, ints) are logged as `rr.Scalars`.
- 3D NumPy arrays that resemble images (e.g., with 1, 3, or 4 channels first) are transposed
from CHW to HWC format and logged as `rr.Image`.
from CHW to HWC format, (optionally) compressed to JPEG and logged as `rr.Image` or `rr.EncodedImage`.
- 1D NumPy arrays are logged as a series of individual scalars, with each element indexed.
- Other multi-dimensional arrays are flattened and logged as individual scalars.
@@ -57,6 +70,7 @@ def log_rerun_data(
Args:
observation: An optional dictionary containing observation data to log.
action: An optional dictionary containing action data to log.
compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality.
"""
if observation:
for k, v in observation.items():
@@ -75,7 +89,8 @@ def log_rerun_data(
for i, vi in enumerate(arr):
rr.log(f"{key}_{i}", rr.Scalars(float(vi)))
else:
rr.log(key, rr.Image(arr), static=True)
img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
rr.log(key, entity=img_entity, static=True)
if action:
for k, v in action.items():