Hugging Face LeRobot

Live-tune a Hugging Face LeRobot robotics policy training run directly from the OmniLoop dashboard.

Hugging Face’s LeRobot is the standard library for robot learning, imitation learning, and reinforcement learning on physical arms and simulators. When training in Hardware-in-the-Loop (HIL) environments (e.g., using SERL or real-arm setups), restarts are extremely costly.

With the OmniLoop LeRobot adapter, you can live-tune learning rates and policy parameters, stream loss metrics, and trigger a freeze-on-exception safety halt when joint limits are hit or control cycles throw errors—allowing you to inspect the robot state, patch parameters, and resume without losing model weights or dropping the arm.

Usage

Instantiate the OmniLoopCallback and wrap your training step loop with the .tick() context manager:

import torch
from lerobot.policies.act.modeling_act import ACTPolicy
from omniloop.integrations.lerobot import OmniLoopCallback

# 1. Initialize your policy, optimizer, and optional scheduler
policy = ACTPolicy(config)
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4)

# 2. Wire up the OmniLoopCallback
callback = OmniLoopCallback(
    policy=policy,
    optimizer=optimizer,
    tunable=["learning_rate"],
    bounds={"learning_rate": (1e-6, 1e-2, 1e-6)}
)

# 3. Wrap your custom training loop
for batch in dataloader:
    with callback.tick():
        # Forward pass & optimization
        output = policy(batch)
        loss = output["loss"]
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # Log loss and other readouts dynamically
        callback.log(loss=loss.item())

Installation

Install the LeRobot extra:

pip install omniloop[lerobot]

How the Adapter Works

Unlike standard reinforcement learning libraries (like Stable-Baselines3), LeRobot does not expose a monolithic callback system because training loops are typically customized per robot setup.

The OmniLoop adapter uses a context-manager pattern (callback.tick()):

• On Enter: Checks the dashboard for a halt command, locks the loop if halted, and applies any pending hyperparameter edits (like writing the new learning_rate to the optimizer’s param_groups).
• On Exit:
  • If a step raises an exception (e.g., joint limit warning, hardware exception, or NaN detection), and freeze_on_exception is enabled, the loop freezes instead of crashing. This gives you a chance to inspect variables, patch safety settings, and resume training.
  • Telemetry is published to the dashboard every publish_every steps.

Configuration

OmniLoopCallback accepts the following parameters:

Parameter	Type	Description
policy	Object	The PyTorch policy model. Attributes on this model can be registered as tunables.
optimizer	Object	The PyTorch optimizer instance to bind learning rate changes to.
scheduler	Object	Optional learning rate scheduler to update alongside optimizer changes.
tunable	List of strings	Hyperparameter names to expose as controls.
bounds	Dictionary	Min, max, and step for sliders {name: (min, max, step)}.
publish_every	Integer	Stream telemetry every N steps.
freeze_on_exception	Boolean	Whether to freeze the run and hold state if an exception occurs during training.
tracker	Object	Optional custom OmniLoop tracker instance to inject.