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:
|
if action_index_before_inference is not None:
|
||||||
indexes_diff = max(0, self.last_index - action_index_before_inference)
|
indexes_diff = max(0, self.last_index - action_index_before_inference)
|
||||||
if indexes_diff != real_delay:
|
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(
|
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,
|
indexes_diff,
|
||||||
real_delay,
|
real_delay,
|
||||||
|
resolved,
|
||||||
)
|
)
|
||||||
return real_delay
|
return resolved
|
||||||
|
|
||||||
return effective_delay
|
return effective_delay
|
||||||
|
|||||||
@@ -116,7 +116,8 @@ def create_inference_engine(
|
|||||||
postprocessor=postprocessor,
|
postprocessor=postprocessor,
|
||||||
robot_wrapper=robot_wrapper,
|
robot_wrapper=robot_wrapper,
|
||||||
rtc_config=config.rtc,
|
rtc_config=config.rtc,
|
||||||
hw_features=hw_features,
|
dataset_features=dataset_features,
|
||||||
|
ordered_action_keys=ordered_action_keys,
|
||||||
task=task,
|
task=task,
|
||||||
fps=fps,
|
fps=fps,
|
||||||
device=device,
|
device=device,
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ import torch
|
|||||||
from lerobot.policies.pretrained import PreTrainedPolicy
|
from lerobot.policies.pretrained import PreTrainedPolicy
|
||||||
from lerobot.policies.rtc import ActionQueue, LatencyTracker, reanchor_relative_rtc_prefix
|
from lerobot.policies.rtc import ActionQueue, LatencyTracker, reanchor_relative_rtc_prefix
|
||||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
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 (
|
from lerobot.processor import (
|
||||||
NormalizerProcessorStep,
|
NormalizerProcessorStep,
|
||||||
PolicyProcessorPipeline,
|
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."""
|
"""Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
|
||||||
if prev_actions.ndim != 2:
|
if prev_actions.ndim != 2:
|
||||||
raise ValueError(f"Expected 2D [T, A] tensor, got shape={tuple(prev_actions.shape)}")
|
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:
|
if steps == target_steps:
|
||||||
return prev_actions
|
return prev_actions
|
||||||
if steps > target_steps:
|
if steps > target_steps:
|
||||||
return prev_actions[:target_steps]
|
return prev_actions[:target_steps]
|
||||||
padded = torch.zeros((target_steps, action_dim), dtype=prev_actions.dtype, device=prev_actions.device)
|
if steps == 0:
|
||||||
padded[:steps] = prev_actions
|
raise ValueError("Cannot pad an empty prefix: no last action to hold.")
|
||||||
return padded
|
# 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,
|
postprocessor: PolicyProcessorPipeline,
|
||||||
robot_wrapper: ThreadSafeRobot,
|
robot_wrapper: ThreadSafeRobot,
|
||||||
rtc_config: RTCConfig,
|
rtc_config: RTCConfig,
|
||||||
hw_features: dict,
|
dataset_features: dict,
|
||||||
|
ordered_action_keys: list[str],
|
||||||
task: str,
|
task: str,
|
||||||
fps: float,
|
fps: float,
|
||||||
device: str | None,
|
device: str | None,
|
||||||
@@ -111,7 +114,8 @@ class RTCInferenceEngine(InferenceEngine):
|
|||||||
self._postprocessor = postprocessor
|
self._postprocessor = postprocessor
|
||||||
self._robot = robot_wrapper
|
self._robot = robot_wrapper
|
||||||
self._rtc_config = rtc_config
|
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._task = task
|
||||||
self._fps = fps
|
self._fps = fps
|
||||||
self._device = device or "cpu"
|
self._device = device or "cpu"
|
||||||
@@ -233,7 +237,13 @@ class RTCInferenceEngine(InferenceEngine):
|
|||||||
"""Pop the next action from the RTC queue (ignores ``obs_frame``)."""
|
"""Pop the next action from the RTC queue (ignores ``obs_frame``)."""
|
||||||
if self._action_queue is None:
|
if self._action_queue is None:
|
||||||
return 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:
|
def notify_observation(self, obs: dict) -> None:
|
||||||
"""Publish the latest observation for the RTC thread to consume."""
|
"""Publish the latest observation for the RTC thread to consume."""
|
||||||
@@ -252,6 +262,8 @@ class RTCInferenceEngine(InferenceEngine):
|
|||||||
policy_device = torch.device(self._device)
|
policy_device = torch.device(self._device)
|
||||||
|
|
||||||
warmup_required = max(1, self._compile_warmup_inferences) if self._use_torch_compile else 0
|
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
|
inference_count = 0
|
||||||
consecutive_errors = 0
|
consecutive_errors = 0
|
||||||
|
|
||||||
@@ -276,7 +288,7 @@ class RTCInferenceEngine(InferenceEngine):
|
|||||||
latency = latency_tracker.max()
|
latency = latency_tracker.max()
|
||||||
delay = math.ceil(latency / time_per_chunk) if latency else 0
|
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 = prepare_observation_for_inference(
|
||||||
obs_batch, policy_device, self._task, self._robot.robot_type
|
obs_batch, policy_device, self._task, self._robot.robot_type
|
||||||
)
|
)
|
||||||
@@ -316,7 +328,8 @@ class RTCInferenceEngine(InferenceEngine):
|
|||||||
inference_count += 1
|
inference_count += 1
|
||||||
consecutive_errors = 0
|
consecutive_errors = 0
|
||||||
is_warmup = self._use_torch_compile and inference_count <= warmup_required
|
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()
|
latency_tracker.reset()
|
||||||
else:
|
else:
|
||||||
latency_tracker.add(new_latency)
|
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,
|
action_index_before_inference=0,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check warning was logged
|
# Check warning was logged (reworded when the discard clamp was added)
|
||||||
assert "Indexes diff is not equal to real delay" in caplog.text
|
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):
|
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
|
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()
|
action_index_before = action_queue_rtc_enabled.get_action_index()
|
||||||
|
for _ in range(5):
|
||||||
|
action_queue_rtc_enabled.get()
|
||||||
|
|
||||||
action_queue_rtc_enabled.merge(
|
action_queue_rtc_enabled.merge(
|
||||||
sample_actions["original"],
|
sample_actions["original"],
|
||||||
|
|||||||
@@ -251,6 +251,69 @@ def test_create_inference_engine_sync():
|
|||||||
assert isinstance(engine, SyncInferenceEngine)
|
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
|
# Pure functions
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
Reference in New Issue
Block a user