mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 03:06:01 +00:00
Compare commits
7 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 6399803b66 | |||
| a19730768f | |||
| 15bc6e0a80 | |||
| eb1c18d172 | |||
| 0441c57356 | |||
| e0b50303aa | |||
| 7845ee6f80 |
@@ -236,11 +236,16 @@ class ActionQueue:
|
||||
if action_index_before_inference is not None:
|
||||
indexes_diff = max(0, self.last_index - action_index_before_inference)
|
||||
if indexes_diff != real_delay:
|
||||
# take the min of both to avoid discarding actions that were not
|
||||
# actually consumed during inference, which would cause a jump in the queue
|
||||
resolved = min(real_delay, indexes_diff)
|
||||
logger.warning(
|
||||
"Indexes diff is not equal to real delay. indexes_diff=%d, real_delay=%d",
|
||||
"Indexes diff != real delay (indexes_diff=%d, real_delay=%d); "
|
||||
"clamping discard to %d to avoid a queue-splice jump.",
|
||||
indexes_diff,
|
||||
real_delay,
|
||||
resolved,
|
||||
)
|
||||
return real_delay
|
||||
return resolved
|
||||
|
||||
return effective_delay
|
||||
|
||||
@@ -116,7 +116,8 @@ def create_inference_engine(
|
||||
postprocessor=postprocessor,
|
||||
robot_wrapper=robot_wrapper,
|
||||
rtc_config=config.rtc,
|
||||
hw_features=hw_features,
|
||||
dataset_features=dataset_features,
|
||||
ordered_action_keys=ordered_action_keys,
|
||||
task=task,
|
||||
fps=fps,
|
||||
device=device,
|
||||
|
||||
@@ -34,7 +34,7 @@ import torch
|
||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||
from lerobot.policies.rtc import ActionQueue, LatencyTracker, reanchor_relative_rtc_prefix
|
||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
||||
from lerobot.policies.utils import prepare_observation_for_inference
|
||||
from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
|
||||
from lerobot.processor import (
|
||||
NormalizerProcessorStep,
|
||||
PolicyProcessorPipeline,
|
||||
@@ -66,14 +66,16 @@ def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int
|
||||
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
|
||||
if prev_actions.ndim != 2:
|
||||
raise ValueError(f"Expected 2D [T, A] tensor, got shape={tuple(prev_actions.shape)}")
|
||||
steps, action_dim = prev_actions.shape
|
||||
steps, _ = prev_actions.shape
|
||||
if steps == target_steps:
|
||||
return prev_actions
|
||||
if steps > target_steps:
|
||||
return prev_actions[:target_steps]
|
||||
padded = torch.zeros((target_steps, action_dim), dtype=prev_actions.dtype, device=prev_actions.device)
|
||||
padded[:steps] = prev_actions
|
||||
return padded
|
||||
if steps == 0:
|
||||
raise ValueError("Cannot pad an empty prefix: no last action to hold.")
|
||||
# repeat the last action to fill the remaining steps to avoid a sudden jump
|
||||
hold = prev_actions[-1:].expand(target_steps - steps, -1)
|
||||
return torch.cat([prev_actions, hold], dim=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -97,7 +99,8 @@ class RTCInferenceEngine(InferenceEngine):
|
||||
postprocessor: PolicyProcessorPipeline,
|
||||
robot_wrapper: ThreadSafeRobot,
|
||||
rtc_config: RTCConfig,
|
||||
hw_features: dict,
|
||||
dataset_features: dict,
|
||||
ordered_action_keys: list[str],
|
||||
task: str,
|
||||
fps: float,
|
||||
device: str | None,
|
||||
@@ -111,7 +114,8 @@ class RTCInferenceEngine(InferenceEngine):
|
||||
self._postprocessor = postprocessor
|
||||
self._robot = robot_wrapper
|
||||
self._rtc_config = rtc_config
|
||||
self._hw_features = hw_features
|
||||
self._obs_features = dataset_features
|
||||
self._ordered_action_keys = ordered_action_keys
|
||||
self._task = task
|
||||
self._fps = fps
|
||||
self._device = device or "cpu"
|
||||
@@ -233,7 +237,13 @@ class RTCInferenceEngine(InferenceEngine):
|
||||
"""Pop the next action from the RTC queue (ignores ``obs_frame``)."""
|
||||
if self._action_queue is None:
|
||||
return None
|
||||
return self._action_queue.get()
|
||||
action = self._action_queue.get()
|
||||
if action is None:
|
||||
return None
|
||||
|
||||
# properly reorder the action dict to match the robot's expected order
|
||||
action_dict = make_robot_action(action, self._obs_features)
|
||||
return torch.tensor([action_dict[k] for k in self._ordered_action_keys])
|
||||
|
||||
def notify_observation(self, obs: dict) -> None:
|
||||
"""Publish the latest observation for the RTC thread to consume."""
|
||||
@@ -252,6 +262,8 @@ class RTCInferenceEngine(InferenceEngine):
|
||||
policy_device = torch.device(self._device)
|
||||
|
||||
warmup_required = max(1, self._compile_warmup_inferences) if self._use_torch_compile else 0
|
||||
# exclude the first N inferences from the latency tracker to avoid cold-start spikes
|
||||
latency_warmup_required = max(1, warmup_required)
|
||||
inference_count = 0
|
||||
consecutive_errors = 0
|
||||
|
||||
@@ -276,7 +288,7 @@ class RTCInferenceEngine(InferenceEngine):
|
||||
latency = latency_tracker.max()
|
||||
delay = math.ceil(latency / time_per_chunk) if latency else 0
|
||||
|
||||
obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation")
|
||||
obs_batch = build_dataset_frame(self._obs_features, obs, prefix="observation")
|
||||
obs_batch = prepare_observation_for_inference(
|
||||
obs_batch, policy_device, self._task, self._robot.robot_type
|
||||
)
|
||||
@@ -316,7 +328,8 @@ class RTCInferenceEngine(InferenceEngine):
|
||||
inference_count += 1
|
||||
consecutive_errors = 0
|
||||
is_warmup = self._use_torch_compile and inference_count <= warmup_required
|
||||
if is_warmup:
|
||||
# Ignore the first N inferences for latency tracking to avoid cold-start spikes
|
||||
if inference_count <= latency_warmup_required:
|
||||
latency_tracker.reset()
|
||||
else:
|
||||
latency_tracker.add(new_latency)
|
||||
|
||||
@@ -457,8 +457,9 @@ def test_merge_validates_delay_consistency(action_queue_rtc_enabled, sample_acti
|
||||
action_index_before_inference=0,
|
||||
)
|
||||
|
||||
# Check warning was logged
|
||||
assert "Indexes diff is not equal to real delay" in caplog.text
|
||||
# Check warning was logged (reworded when the discard clamp was added)
|
||||
assert "Indexes diff != real delay" in caplog.text
|
||||
assert "clamping discard" in caplog.text
|
||||
|
||||
|
||||
def test_merge_no_warning_when_delays_match(action_queue_rtc_enabled, sample_actions, caplog):
|
||||
@@ -790,8 +791,12 @@ def test_typical_rtc_workflow(action_queue_rtc_enabled, sample_actions):
|
||||
|
||||
assert action_queue_rtc_enabled.qsize() == 40
|
||||
|
||||
# Second inference with delay
|
||||
# Second inference with delay. Capture the index at inference *start*, then consume
|
||||
# `real_delay` actions to simulate the robot executing during inference, so the
|
||||
# measured `indexes_diff` matches `real_delay` and the discard clamp is a no-op.
|
||||
action_index_before = action_queue_rtc_enabled.get_action_index()
|
||||
for _ in range(5):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
action_queue_rtc_enabled.merge(
|
||||
sample_actions["original"],
|
||||
|
||||
@@ -251,6 +251,69 @@ def test_create_inference_engine_sync():
|
||||
assert isinstance(engine, SyncInferenceEngine)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RTC action-key ordering and prefix padding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_rtc_get_action_remaps_model_order_to_ordered_action_keys():
|
||||
"""RTC must remap the model-order action vector to ``ordered_action_keys`` by NAME
|
||||
before returning — matching sync. Otherwise the strategy maps model outputs onto the
|
||||
wrong joints (a per-joint permutation) whenever the two orders differ."""
|
||||
from lerobot.rollout import RTCInferenceConfig, RTCInferenceEngine, create_inference_engine
|
||||
from lerobot.utils.constants import ACTION
|
||||
|
||||
# The model emits actions in dataset order [a, b, c]; the robot wants [c, a, b].
|
||||
dataset_action_names = ["a.pos", "b.pos", "c.pos"]
|
||||
ordered_action_keys = ["c.pos", "a.pos", "b.pos"]
|
||||
dataset_features = {
|
||||
ACTION: {"dtype": "float32", "shape": (3,), "names": dataset_action_names},
|
||||
}
|
||||
|
||||
engine = create_inference_engine(
|
||||
RTCInferenceConfig(),
|
||||
policy=MagicMock(),
|
||||
preprocessor=MagicMock(steps=[]),
|
||||
postprocessor=MagicMock(steps=[]),
|
||||
robot_wrapper=MagicMock(robot_type="mock"),
|
||||
hw_features={},
|
||||
dataset_features=dataset_features,
|
||||
ordered_action_keys=ordered_action_keys,
|
||||
task="test",
|
||||
fps=30.0,
|
||||
device="cpu",
|
||||
)
|
||||
assert isinstance(engine, RTCInferenceEngine)
|
||||
|
||||
# Queue yields the model-order vector a=1, b=2, c=3.
|
||||
engine._action_queue = MagicMock()
|
||||
engine._action_queue.get.return_value = torch.tensor([1.0, 2.0, 3.0])
|
||||
|
||||
out = engine.get_action(None)
|
||||
# Remapped by name to [c, a, b] = [3, 1, 2]; positional pass-through would give [1, 2, 3].
|
||||
torch.testing.assert_close(out, torch.tensor([3.0, 1.0, 2.0]))
|
||||
|
||||
|
||||
def test_normalize_prev_actions_length_holds_last_action_not_zeros():
|
||||
"""A short RTC prefix is padded by repeating the last action, never with zeros —
|
||||
zeros decode to the mean action and cause the intermittent chunk-seam spike."""
|
||||
from lerobot.rollout.inference.rtc import _normalize_prev_actions_length
|
||||
|
||||
prev = torch.tensor([[1.0, -1.0], [2.0, -2.0]]) # 2 steps, dim 2
|
||||
|
||||
# Pad up to 5: rows 2..4 must equal the last real row, not zeros.
|
||||
padded = _normalize_prev_actions_length(prev, target_steps=5)
|
||||
assert padded.shape == (5, 2)
|
||||
torch.testing.assert_close(padded[:2], prev)
|
||||
for i in range(2, 5):
|
||||
torch.testing.assert_close(padded[i], prev[-1])
|
||||
assert not torch.any(padded[2:] == 0.0), "pad rows must not be zeros"
|
||||
|
||||
# Exact length: unchanged. Truncation: first `target_steps` rows.
|
||||
torch.testing.assert_close(_normalize_prev_actions_length(prev, 2), prev)
|
||||
torch.testing.assert_close(_normalize_prev_actions_length(prev, 1), prev[:1])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user