mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-22 12:09:42 +00:00
[RTC] Real Time Chunking for Pi0, Smolvla, Pi0.5 (#1698)
* Add Real-Time Chunking (RTC) support for flow matching models Implement Real-Time Chunking (RTC) for action chunking policies using flow matching denoising. RTC enables smooth action transitions between consecutive chunks by using prefix guidance during denoising. Key features: - RTCProcessor class with denoise_step method for RTC guidance - Tracker system for debug tracking using time-based dictionary storage - RTCDebugVisualizer with comprehensive visualization utilities - Integration with SmolVLA policy for flow matching models - Support for multiple prefix attention schedules (ZEROS, ONES, LINEAR, EXP) - Configurable execution horizon and max guidance weight - Example scripts for dataset evaluation and real-time control Technical details: - Uses autograd-based gradient computation for RTC corrections - Time-based tracking eliminates duplicate step issues - Proxy methods in RTCProcessor for cleaner API - Full integration with LeRobot's policy and dataset systems Files added/modified: - src/lerobot/configs/types.py: Add RTCAttentionSchedule enum - src/lerobot/policies/rtc/: Core RTC implementation - configuration_rtc.py: RTC configuration - modeling_rtc.py: RTCProcessor with denoise_step - debug_handler.py: Tracker for debug information - debug_visualizer.py: Visualization utilities - src/lerobot/policies/smolvla/modeling_smolvla.py: RTC integration - examples/rtc/: Example scripts and evaluation tools 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> Co-Authored-By: Claude <noreply@anthropic.com> * Fix rtc_config attribute access in SmolVLA Use getattr() to safely check for rtc_config attribute existence instead of direct attribute access. This fixes AttributeError when loading policies without rtc_config in their config. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> Co-Authored-By: Claude <noreply@anthropic.com> * fixup! Fix rtc_config attribute access in SmolVLA * Add RTCConfig field to SmolVLAConfig Add rtc_config as an optional field in SmolVLAConfig to properly support Real-Time Chunking configuration. This replaces the previous getattr() workarounds with direct attribute access, making the code cleaner and more maintainable. Changes: - Import RTCConfig in configuration_smolvla.py - Add rtc_config: RTCConfig | None = None field - Revert getattr() calls to direct attribute access in modeling_smolvla.py 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> Co-Authored-By: Claude <noreply@anthropic.com> * Refactor RTC enabled checks to use _rtc_enabled helper Add _rtc_enabled() helper method in VLAFlowMatching class to simplify and clean up RTC enabled checks throughout the code. This reduces code duplication and improves readability. Changes: - Add _rtc_enabled() method in VLAFlowMatching - Replace verbose rtc_config checks with _rtc_enabled() calls - Maintain exact same functionality with cleaner code 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> Co-Authored-By: Claude <noreply@anthropic.com> * Rename track_debug method to track Simplify the method name from track_debug to just track for better readability and consistency. The method already has clear documentation about its debug tracking purpose. Changes: - Rename RTCProcessor.track_debug() to track() - Update all call sites in modeling_smolvla.py and modeling_rtc.py 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> Co-Authored-By: Claude <noreply@anthropic.com> * Use output_dir for saving all evaluation images Update eval_dataset.py to save all comparison images to the configured output_dir instead of the current directory. This provides better organization and allows users to specify where outputs should be saved. Changes: - Add os import at top level - Create output_dir at start of run_evaluation() - Save all comparison images to output_dir - Remove duplicate os imports - Update init_rtc_processor() docstring to be more concise 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> Co-Authored-By: Claude <noreply@anthropic.com> * fixup! Use output_dir for saving all evaluation images * Fix logging buffering and enable tracking when RTC config provided - Add force=True to logging.basicConfig to override existing configuration - Enable line buffering for stdout/stderr for real-time log output - Modify init_rtc_processor to create processor when rtc_config exists even if RTC is disabled, allowing tracking of denoising data 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> * Refactor SmolVLA plotting to use tracker data instead of local variables Remove local tracking variables (correction, x1_t, error) from the denoising loop and instead retrieve plotting data from the RTC tracker after each denoise step. This makes the code cleaner and uses the tracker as the single source of truth for debug/visualization data. Changes: - Remove initialization of correction, x1_t, error before denoising loop - After each Euler step, retrieve most recent debug step from tracker - Extract correction, x1_t, err from debug step for plotting - Update tracking condition to use is_debug_enabled() method 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> * Move plotting logic from modeling_smolvla to eval_dataset script Refactor to improve separation of concerns: modeling_smolvla.py changes: - Remove all plotting logic from sample_actions method - Remove viz_xt_axs, viz_vt_axs, viz_x1t_axs parameters - Remove matplotlib and RTCDebugVisualizer imports - Remove viz_fig, viz_axs, denoise_step_counter instance variables - Simplify denoising loop to only track data in rtc_processor eval_dataset.py changes: - Add _plot_denoising_steps_from_tracker helper method - Retrieve debug steps from tracker after inference - Plot x_t, v_t, x1_t, correction, and error from tracker data - Enable debug tracking (cfg.rtc.debug = True) for visualization - Remove viz axes parameters from predict_action_chunk calls modeling_rtc.py changes: - Remove v_t from track() call (handled by user change) Benefits: - Cleaner modeling code focused on inference - Evaluation script owns all visualization logic - Better separation of concerns - Tracker is single source of truth for debug data 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> * Refactor plotting loging * fixup! Refactor plotting loging * Improve visualization: separate correction plot and fix axis scaling Changes: - Create separate figure for correction data instead of overlaying on v_t - Add _rescale_axes helper method to properly scale all axes - Add 10% margin to y-axis for better visualization - Fix v_t chart vertical compression issue Benefits: - Clearer v_t plot without correction overlay - Better axis scaling with proper margins - Separate correction figure for focused analysis - Improved readability of all denoising visualizations Output files: - denoising_xt_comparison.png (x_t trajectories) - denoising_vt_comparison.png (v_t velocity - now cleaner) - denoising_correction_comparison.png (NEW - separate corrections) - denoising_x1t_comparison.png (x1_t state with error) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com> * fixup! Improve visualization: separate correction plot and fix axis scaling * fixup! fixup! Improve visualization: separate correction plot and fix axis scaling * fixup! fixup! fixup! Improve visualization: separate correction plot and fix axis scaling * Fix traacking * Right kwargs for the policy * Add tests for tracker * Fix tests * Drop not required methods * Add torch compilation for eval_dataset * delete policies * Add matplotliv to dev * fixup! Add matplotliv to dev * Experiemnt with late detach * Debug * Fix compilation * Add RTC to PI0 * Pi0 * Pi0 eval dataset * fixup! Pi0 eval dataset * Turn off compilation for pi0/pi05 * fixup! Turn off compilation for pi0/pi05 * fixup! fixup! Turn off compilation for pi0/pi05 * fixup! fixup! fixup! Turn off compilation for pi0/pi05 * fixup! fixup! fixup! fixup! Turn off compilation for pi0/pi05 * fixup! fixup! fixup! fixup! fixup! Turn off compilation for pi0/pi05 * Add workable flow * Small fixes * Add more tests * Add validatio at the end * Update README * Silent validation * Fix tests * Add tests for modeling_rtc * Add tests for flow matching models with RTC * fixup! Add tests for flow matching models with RTC * fixup! fixup! Add tests for flow matching models with RTC * Add one more test * fixup! Add one more test * Fix test to use _rtc_enabled() instead of is_rtc_enabled() 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fixup! Fix test to use _rtc_enabled() instead of is_rtc_enabled() * fixup! fixup! Fix test to use _rtc_enabled() instead of is_rtc_enabled() * Add RTC initialization tests without config for PI0.5 and SmolVLA Add test_pi05_rtc_initialization_without_rtc_config and test_smolvla_rtc_initialization_without_rtc_config to verify that policies can initialize without RTC config and that _rtc_enabled() returns False in this case. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix PI0.5 init_rtc_processor to use getattr instead of direct model access 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix SmolVLA init_rtc_processor to use getattr instead of direct model access 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * Fix PI0.5 RTC tests to use quantile stats (q01, q99) for normalization 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com> * fixup! Fix PI0.5 RTC tests to use quantile stats (q01, q99) for normalization * Fixup eval with real robot * fixup! Fixup eval with real robot * fixup! fixup! Fixup eval with real robot * Extract simulator logic from eval_with real robot and add proper headers to files * Update images * Fix tests * fixup! Fix tests * add docs for rtc * enhance doc and add images * Fix instal instructions --------- Co-authored-by: Ben Zhang <benzhangniu@gmail.com> Co-authored-by: Alexander Soare <alexander.soare159@gmail.com> Co-authored-by: Michel Aractingi <michel.aractingi@huggingface.co>
This commit is contained in:
Eugene Mironov
@@ -0,0 +1,825 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for RTC ActionQueue module."""
|
||||
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.rtc.action_queue import ActionQueue
|
||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
||||
|
||||
# ====================== Fixtures ======================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtc_config_enabled():
|
||||
"""Create an RTC config with RTC enabled."""
|
||||
return RTCConfig(enabled=True, execution_horizon=10, max_guidance_weight=1.0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtc_config_disabled():
|
||||
"""Create an RTC config with RTC disabled."""
|
||||
return RTCConfig(enabled=False, execution_horizon=10, max_guidance_weight=1.0)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_actions():
|
||||
"""Create sample action tensors for testing."""
|
||||
return {
|
||||
"original": torch.randn(50, 6), # (time_steps, action_dim)
|
||||
"processed": torch.randn(50, 6),
|
||||
"short": torch.randn(10, 6),
|
||||
"longer": torch.randn(100, 6),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def action_queue_rtc_enabled(rtc_config_enabled):
|
||||
"""Create an ActionQueue with RTC enabled."""
|
||||
return ActionQueue(rtc_config_enabled)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def action_queue_rtc_disabled(rtc_config_disabled):
|
||||
"""Create an ActionQueue with RTC disabled."""
|
||||
return ActionQueue(rtc_config_disabled)
|
||||
|
||||
|
||||
# ====================== Initialization Tests ======================
|
||||
|
||||
|
||||
def test_action_queue_initialization_rtc_enabled(rtc_config_enabled):
|
||||
"""Test ActionQueue initializes correctly with RTC enabled."""
|
||||
queue = ActionQueue(rtc_config_enabled)
|
||||
assert queue.queue is None
|
||||
assert queue.original_queue is None
|
||||
assert queue.last_index == 0
|
||||
assert queue.cfg.enabled is True
|
||||
|
||||
|
||||
def test_action_queue_initialization_rtc_disabled(rtc_config_disabled):
|
||||
"""Test ActionQueue initializes correctly with RTC disabled."""
|
||||
queue = ActionQueue(rtc_config_disabled)
|
||||
assert queue.queue is None
|
||||
assert queue.original_queue is None
|
||||
assert queue.last_index == 0
|
||||
assert queue.cfg.enabled is False
|
||||
|
||||
|
||||
# ====================== get() Tests ======================
|
||||
|
||||
|
||||
def test_get_returns_none_when_empty(action_queue_rtc_enabled):
|
||||
"""Test get() returns None when queue is empty."""
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert action is None
|
||||
|
||||
|
||||
def test_get_returns_actions_sequentially(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get() returns actions in sequence."""
|
||||
# Initialize queue with actions
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
|
||||
# Get first action
|
||||
action1 = action_queue_rtc_enabled.get()
|
||||
assert action1 is not None
|
||||
assert action1.shape == (6,)
|
||||
assert torch.equal(action1, sample_actions["processed"][0])
|
||||
|
||||
# Get second action
|
||||
action2 = action_queue_rtc_enabled.get()
|
||||
assert action2 is not None
|
||||
assert torch.equal(action2, sample_actions["processed"][1])
|
||||
|
||||
|
||||
def test_get_returns_none_after_exhaustion(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get() returns None after all actions are consumed."""
|
||||
# Use short action sequence
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume all actions
|
||||
for _ in range(10):
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert action is not None
|
||||
|
||||
# Next get should return None
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert action is None
|
||||
|
||||
|
||||
def test_get_increments_last_index(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get() increments last_index correctly."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
|
||||
assert action_queue_rtc_enabled.last_index == 0
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.last_index == 1
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.last_index == 2
|
||||
|
||||
|
||||
# ====================== qsize() Tests ======================
|
||||
|
||||
|
||||
def test_qsize_returns_zero_when_empty(action_queue_rtc_enabled):
|
||||
"""Test qsize() returns 0 when queue is empty."""
|
||||
assert action_queue_rtc_enabled.qsize() == 0
|
||||
|
||||
|
||||
def test_qsize_returns_correct_size(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test qsize() returns correct number of remaining actions."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
assert action_queue_rtc_enabled.qsize() == 10
|
||||
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.qsize() == 9
|
||||
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.qsize() == 8
|
||||
|
||||
|
||||
def test_qsize_after_exhaustion(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test qsize() returns 0 after queue is exhausted."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume all actions
|
||||
for _ in range(10):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
assert action_queue_rtc_enabled.qsize() == 0
|
||||
|
||||
|
||||
# ====================== empty() Tests ======================
|
||||
|
||||
|
||||
def test_empty_returns_true_when_empty(action_queue_rtc_enabled):
|
||||
"""Test empty() returns True when queue is empty."""
|
||||
assert action_queue_rtc_enabled.empty() is True
|
||||
|
||||
|
||||
def test_empty_returns_false_when_not_empty(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test empty() returns False when queue has actions."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
assert action_queue_rtc_enabled.empty() is False
|
||||
|
||||
|
||||
def test_empty_after_partial_consumption(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test empty() returns False after partial consumption."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
action_queue_rtc_enabled.get()
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
assert action_queue_rtc_enabled.empty() is False
|
||||
|
||||
|
||||
def test_empty_after_full_consumption(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test empty() returns True after all actions consumed."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume all
|
||||
for _ in range(10):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
assert action_queue_rtc_enabled.empty() is True
|
||||
|
||||
|
||||
# ====================== get_action_index() Tests ======================
|
||||
|
||||
|
||||
def test_get_action_index_initial_value(action_queue_rtc_enabled):
|
||||
"""Test get_action_index() returns 0 initially."""
|
||||
assert action_queue_rtc_enabled.get_action_index() == 0
|
||||
|
||||
|
||||
def test_get_action_index_after_consumption(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get_action_index() tracks consumption correctly."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
assert action_queue_rtc_enabled.get_action_index() == 0
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.get_action_index() == 1
|
||||
action_queue_rtc_enabled.get()
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.get_action_index() == 3
|
||||
|
||||
|
||||
# ====================== get_left_over() Tests ======================
|
||||
|
||||
|
||||
def test_get_left_over_returns_none_when_empty(action_queue_rtc_enabled):
|
||||
"""Test get_left_over() returns None when queue is empty."""
|
||||
leftover = action_queue_rtc_enabled.get_left_over()
|
||||
assert leftover is None
|
||||
|
||||
|
||||
def test_get_left_over_returns_all_when_unconsumed(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get_left_over() returns all original actions when none consumed."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
leftover = action_queue_rtc_enabled.get_left_over()
|
||||
assert leftover is not None
|
||||
assert leftover.shape == (10, 6)
|
||||
assert torch.equal(leftover, sample_actions["short"])
|
||||
|
||||
|
||||
def test_get_left_over_returns_remaining_after_consumption(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get_left_over() returns only remaining original actions."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume 3 actions
|
||||
action_queue_rtc_enabled.get()
|
||||
action_queue_rtc_enabled.get()
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
leftover = action_queue_rtc_enabled.get_left_over()
|
||||
assert leftover is not None
|
||||
assert leftover.shape == (7, 6)
|
||||
assert torch.equal(leftover, sample_actions["short"][3:])
|
||||
|
||||
|
||||
def test_get_left_over_returns_empty_after_exhaustion(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get_left_over() returns empty tensor after all consumed."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume all
|
||||
for _ in range(10):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
leftover = action_queue_rtc_enabled.get_left_over()
|
||||
assert leftover is not None
|
||||
assert leftover.shape == (0, 6)
|
||||
|
||||
|
||||
# ====================== merge() with RTC Enabled Tests ======================
|
||||
|
||||
|
||||
def test_merge_replaces_queue_when_rtc_enabled(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test merge() replaces queue when RTC is enabled."""
|
||||
# Add initial actions
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
assert action_queue_rtc_enabled.qsize() == 10
|
||||
|
||||
# Consume some actions
|
||||
action_queue_rtc_enabled.get()
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.qsize() == 8
|
||||
|
||||
# Merge new actions - should replace, not append
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=5)
|
||||
|
||||
# Queue should be replaced with new actions minus delay
|
||||
# Original has 50 actions, delay is 5, so remaining is 45
|
||||
assert action_queue_rtc_enabled.qsize() == 45
|
||||
assert action_queue_rtc_enabled.get_action_index() == 0
|
||||
|
||||
|
||||
def test_merge_respects_real_delay(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test merge() correctly applies real_delay when RTC is enabled."""
|
||||
delay = 10
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=delay)
|
||||
|
||||
# Queue should have original length minus delay
|
||||
expected_size = len(sample_actions["original"]) - delay
|
||||
assert action_queue_rtc_enabled.qsize() == expected_size
|
||||
|
||||
# First action should be the one at index [delay]
|
||||
first_action = action_queue_rtc_enabled.get()
|
||||
assert torch.equal(first_action, sample_actions["processed"][delay])
|
||||
|
||||
|
||||
def test_merge_resets_last_index_when_rtc_enabled(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test merge() resets last_index to 0 when RTC is enabled."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
action_queue_rtc_enabled.get()
|
||||
action_queue_rtc_enabled.get()
|
||||
assert action_queue_rtc_enabled.last_index == 2
|
||||
|
||||
# Merge new actions
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=5)
|
||||
|
||||
assert action_queue_rtc_enabled.last_index == 0
|
||||
|
||||
|
||||
def test_merge_with_zero_delay(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test merge() with zero delay keeps all actions."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
|
||||
assert action_queue_rtc_enabled.qsize() == len(sample_actions["original"])
|
||||
|
||||
|
||||
def test_merge_with_large_delay(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test merge() with delay larger than action sequence."""
|
||||
# Delay is larger than sequence length
|
||||
delay = 100
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=delay)
|
||||
|
||||
# Queue should be empty (delay >= length)
|
||||
assert action_queue_rtc_enabled.qsize() == 0
|
||||
|
||||
|
||||
# ====================== merge() with RTC Disabled Tests ======================
|
||||
|
||||
|
||||
def test_merge_appends_when_rtc_disabled(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test merge() appends actions when RTC is disabled."""
|
||||
# Add initial actions
|
||||
action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
initial_size = action_queue_rtc_disabled.qsize()
|
||||
assert initial_size == 10
|
||||
|
||||
# Merge more actions
|
||||
action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Should have appended
|
||||
assert action_queue_rtc_disabled.qsize() == initial_size + 10
|
||||
|
||||
|
||||
def test_merge_removes_consumed_actions_when_appending(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test merge() removes consumed actions before appending when RTC is disabled."""
|
||||
# Add initial actions
|
||||
action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
assert action_queue_rtc_disabled.qsize() == 10
|
||||
|
||||
# Consume 3 actions
|
||||
action_queue_rtc_disabled.get()
|
||||
action_queue_rtc_disabled.get()
|
||||
action_queue_rtc_disabled.get()
|
||||
assert action_queue_rtc_disabled.qsize() == 7
|
||||
|
||||
# Merge more actions
|
||||
action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Should have 7 remaining + 10 new = 17
|
||||
assert action_queue_rtc_disabled.qsize() == 17
|
||||
|
||||
|
||||
def test_merge_resets_last_index_after_append(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test merge() resets last_index after appending when RTC is disabled."""
|
||||
action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
action_queue_rtc_disabled.get()
|
||||
action_queue_rtc_disabled.get()
|
||||
assert action_queue_rtc_disabled.last_index == 2
|
||||
|
||||
# Merge more actions
|
||||
action_queue_rtc_disabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# last_index should be reset to 0
|
||||
assert action_queue_rtc_disabled.last_index == 0
|
||||
|
||||
|
||||
def test_merge_ignores_delay_when_rtc_disabled(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test merge() ignores real_delay parameter when RTC is disabled."""
|
||||
action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=10)
|
||||
|
||||
# All actions should be in queue (delay ignored)
|
||||
assert action_queue_rtc_disabled.qsize() == len(sample_actions["original"])
|
||||
|
||||
|
||||
def test_merge_first_call_with_rtc_disabled(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test merge() on first call with RTC disabled."""
|
||||
action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
|
||||
assert action_queue_rtc_disabled.qsize() == len(sample_actions["original"])
|
||||
assert action_queue_rtc_disabled.last_index == 0
|
||||
|
||||
|
||||
# ====================== merge() with Different Action Shapes Tests ======================
|
||||
|
||||
|
||||
def test_merge_with_different_action_dims():
|
||||
"""Test merge() handles actions with different dimensions."""
|
||||
cfg = RTCConfig(enabled=True, execution_horizon=10)
|
||||
queue = ActionQueue(cfg)
|
||||
|
||||
# Actions with 4 dimensions instead of 6
|
||||
actions_4d = torch.randn(20, 4)
|
||||
queue.merge(actions_4d, actions_4d, real_delay=5)
|
||||
|
||||
action = queue.get()
|
||||
assert action.shape == (4,)
|
||||
|
||||
|
||||
def test_merge_with_different_lengths():
|
||||
"""Test merge() handles action sequences of varying lengths."""
|
||||
cfg = RTCConfig(enabled=False, execution_horizon=10)
|
||||
queue = ActionQueue(cfg)
|
||||
|
||||
# Add sequences of different lengths
|
||||
queue.merge(torch.randn(10, 6), torch.randn(10, 6), real_delay=0)
|
||||
assert queue.qsize() == 10
|
||||
|
||||
queue.merge(torch.randn(25, 6), torch.randn(25, 6), real_delay=0)
|
||||
assert queue.qsize() == 35
|
||||
|
||||
|
||||
# ====================== merge() Delay Validation Tests ======================
|
||||
|
||||
|
||||
def test_merge_validates_delay_consistency(action_queue_rtc_enabled, sample_actions, caplog):
|
||||
"""Test merge() validates that real_delay matches action index difference."""
|
||||
import logging
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
# Initialize queue
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume 5 actions
|
||||
for _ in range(5):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
# Merge with mismatched delay (should log warning)
|
||||
# We consumed 5 actions, so index is 5. If we pass action_index_before_inference=0,
|
||||
# then indexes_diff=5, but if real_delay=3, it will warn
|
||||
action_queue_rtc_enabled.merge(
|
||||
sample_actions["original"],
|
||||
sample_actions["processed"],
|
||||
real_delay=3,
|
||||
action_index_before_inference=0,
|
||||
)
|
||||
|
||||
# Check warning was logged
|
||||
assert "Indexes diff is not equal to real delay" in caplog.text
|
||||
|
||||
|
||||
def test_merge_no_warning_when_delays_match(action_queue_rtc_enabled, sample_actions, caplog):
|
||||
"""Test merge() doesn't warn when delays are consistent."""
|
||||
import logging
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
# Initialize queue
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume 5 actions
|
||||
for _ in range(5):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
# Merge with matching delay
|
||||
action_queue_rtc_enabled.merge(
|
||||
sample_actions["original"],
|
||||
sample_actions["processed"],
|
||||
real_delay=5,
|
||||
action_index_before_inference=0,
|
||||
)
|
||||
|
||||
# Should not have warning
|
||||
assert "Indexes diff is not equal to real delay" not in caplog.text
|
||||
|
||||
|
||||
def test_merge_skips_validation_when_action_index_none(action_queue_rtc_enabled, sample_actions, caplog):
|
||||
"""Test merge() skips delay validation when action_index_before_inference is None."""
|
||||
import logging
|
||||
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
for _ in range(5):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
# Pass None for action_index_before_inference
|
||||
action_queue_rtc_enabled.merge(
|
||||
sample_actions["original"],
|
||||
sample_actions["processed"],
|
||||
real_delay=999, # Doesn't matter
|
||||
action_index_before_inference=None,
|
||||
)
|
||||
|
||||
# Should not warn (validation skipped)
|
||||
assert "Indexes diff is not equal to real delay" not in caplog.text
|
||||
|
||||
|
||||
# ====================== Thread Safety Tests ======================
|
||||
|
||||
|
||||
def test_get_is_thread_safe(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get() is thread-safe with multiple consumers."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["longer"], sample_actions["longer"], real_delay=0)
|
||||
|
||||
results = []
|
||||
errors = []
|
||||
|
||||
def consumer():
|
||||
try:
|
||||
for _ in range(25):
|
||||
action = action_queue_rtc_enabled.get()
|
||||
if action is not None:
|
||||
results.append(action)
|
||||
time.sleep(0.001)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=consumer) for _ in range(4)]
|
||||
|
||||
for t in threads:
|
||||
t.start()
|
||||
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Should not have errors
|
||||
assert len(errors) == 0
|
||||
|
||||
# Should have consumed all actions (100 total, 4 threads * 25 each)
|
||||
assert len(results) == 100
|
||||
|
||||
# All results should be unique (no duplicate consumption)
|
||||
# We can verify by checking that indices are not duplicated
|
||||
# Since we don't track indices in results, we check total count is correct
|
||||
assert action_queue_rtc_enabled.qsize() == 0
|
||||
|
||||
|
||||
def test_merge_is_thread_safe(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test merge() is thread-safe with multiple producers."""
|
||||
errors = []
|
||||
|
||||
def producer():
|
||||
try:
|
||||
for _ in range(5):
|
||||
action_queue_rtc_disabled.merge(
|
||||
sample_actions["short"], sample_actions["short"], real_delay=0
|
||||
)
|
||||
time.sleep(0.001)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=producer) for _ in range(3)]
|
||||
|
||||
for t in threads:
|
||||
t.start()
|
||||
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
# Should not have errors
|
||||
assert len(errors) == 0
|
||||
|
||||
# Should have accumulated all actions (3 threads * 5 merges * 10 actions = 150)
|
||||
assert action_queue_rtc_disabled.qsize() == 150
|
||||
|
||||
|
||||
def test_concurrent_get_and_merge(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test concurrent get() and merge() operations."""
|
||||
errors = []
|
||||
consumed_count = [0]
|
||||
|
||||
def consumer():
|
||||
try:
|
||||
for _ in range(50):
|
||||
action = action_queue_rtc_disabled.get()
|
||||
if action is not None:
|
||||
consumed_count[0] += 1
|
||||
time.sleep(0.001)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
def producer():
|
||||
try:
|
||||
for _ in range(10):
|
||||
action_queue_rtc_disabled.merge(
|
||||
sample_actions["short"], sample_actions["short"], real_delay=0
|
||||
)
|
||||
time.sleep(0.005)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
consumer_threads = [threading.Thread(target=consumer) for _ in range(2)]
|
||||
producer_threads = [threading.Thread(target=producer) for _ in range(2)]
|
||||
|
||||
for t in consumer_threads + producer_threads:
|
||||
t.start()
|
||||
|
||||
for t in consumer_threads + producer_threads:
|
||||
t.join()
|
||||
|
||||
# Should not have errors
|
||||
assert len(errors) == 0
|
||||
|
||||
# Should have consumed some or all actions (non-deterministic due to timing)
|
||||
# Total produced: 2 producers * 10 merges * 10 actions = 200
|
||||
# Total consumed attempts: 2 consumers * 50 = 100
|
||||
assert consumed_count[0] <= 200
|
||||
|
||||
|
||||
# ====================== get_left_over() Thread Safety Tests ======================
|
||||
|
||||
|
||||
def test_get_left_over_is_thread_safe(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test get_left_over() is thread-safe with concurrent access."""
|
||||
action_queue_rtc_enabled.merge(sample_actions["longer"], sample_actions["longer"], real_delay=0)
|
||||
|
||||
errors = []
|
||||
leftovers = []
|
||||
|
||||
def reader():
|
||||
try:
|
||||
for _ in range(20):
|
||||
leftover = action_queue_rtc_enabled.get_left_over()
|
||||
if leftover is not None:
|
||||
leftovers.append(leftover.shape[0])
|
||||
time.sleep(0.001)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=reader) for _ in range(3)]
|
||||
|
||||
# Also consume some actions concurrently
|
||||
def consumer():
|
||||
try:
|
||||
for _ in range(10):
|
||||
action_queue_rtc_enabled.get()
|
||||
time.sleep(0.002)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
consumer_thread = threading.Thread(target=consumer)
|
||||
|
||||
all_threads = threads + [consumer_thread]
|
||||
|
||||
for t in all_threads:
|
||||
t.start()
|
||||
|
||||
for t in all_threads:
|
||||
t.join()
|
||||
|
||||
# Should not have errors
|
||||
assert len(errors) == 0
|
||||
|
||||
# Leftovers should be monotonically decreasing or stable
|
||||
# (as actions are consumed, leftover size decreases)
|
||||
assert len(leftovers) > 0
|
||||
|
||||
|
||||
# ====================== Edge Cases Tests ======================
|
||||
|
||||
|
||||
def test_queue_with_single_action(action_queue_rtc_enabled):
|
||||
"""Test queue behavior with a single action."""
|
||||
single_action_original = torch.randn(1, 6)
|
||||
single_action_processed = torch.randn(1, 6)
|
||||
|
||||
action_queue_rtc_enabled.merge(single_action_original, single_action_processed, real_delay=0)
|
||||
|
||||
assert action_queue_rtc_enabled.qsize() == 1
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert action is not None
|
||||
assert action.shape == (6,)
|
||||
assert action_queue_rtc_enabled.qsize() == 0
|
||||
|
||||
|
||||
def test_queue_behavior_after_multiple_merge_cycles(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test queue maintains correct state through multiple merge cycles."""
|
||||
for _ in range(5):
|
||||
action_queue_rtc_enabled.merge(sample_actions["short"], sample_actions["short"], real_delay=0)
|
||||
|
||||
# Consume half
|
||||
for _ in range(5):
|
||||
action_queue_rtc_enabled.get()
|
||||
|
||||
# Merge again
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=3)
|
||||
|
||||
assert action_queue_rtc_enabled.qsize() > 0
|
||||
|
||||
|
||||
def test_queue_with_all_zeros_actions(action_queue_rtc_enabled):
|
||||
"""Test queue handles all-zero action tensors."""
|
||||
zeros_actions = torch.zeros(20, 6)
|
||||
action_queue_rtc_enabled.merge(zeros_actions, zeros_actions, real_delay=0)
|
||||
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert torch.all(action == 0)
|
||||
|
||||
|
||||
def test_queue_clones_input_tensors(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test that merge() clones input tensors, not storing references."""
|
||||
original_copy = sample_actions["original"].clone()
|
||||
processed_copy = sample_actions["processed"].clone()
|
||||
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
|
||||
# Modify original tensors
|
||||
sample_actions["original"].fill_(999.0)
|
||||
sample_actions["processed"].fill_(-999.0)
|
||||
|
||||
# Queue should have cloned values
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert not torch.equal(action, sample_actions["processed"][0])
|
||||
assert torch.equal(action, processed_copy[0])
|
||||
|
||||
leftover = action_queue_rtc_enabled.get_left_over()
|
||||
assert not torch.equal(leftover, sample_actions["original"][1:])
|
||||
assert torch.equal(leftover, original_copy[1:])
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_queue_handles_gpu_tensors():
|
||||
"""Test queue correctly handles GPU tensors."""
|
||||
cfg = RTCConfig(enabled=True, execution_horizon=10)
|
||||
queue = ActionQueue(cfg)
|
||||
|
||||
actions_gpu = torch.randn(20, 6, device="cuda")
|
||||
queue.merge(actions_gpu, actions_gpu, real_delay=0)
|
||||
|
||||
action = queue.get()
|
||||
assert action.device.type == "cuda"
|
||||
|
||||
leftover = queue.get_left_over()
|
||||
assert leftover.device.type == "cuda"
|
||||
|
||||
|
||||
def test_queue_handles_different_dtypes():
|
||||
"""Test queue handles actions with different dtypes."""
|
||||
cfg = RTCConfig(enabled=True, execution_horizon=10)
|
||||
queue = ActionQueue(cfg)
|
||||
|
||||
# Use float64 instead of default float32
|
||||
actions_f64 = torch.randn(20, 6, dtype=torch.float64)
|
||||
queue.merge(actions_f64, actions_f64, real_delay=0)
|
||||
|
||||
action = queue.get()
|
||||
assert action.dtype == torch.float64
|
||||
|
||||
|
||||
def test_empty_with_none_queue(action_queue_rtc_enabled):
|
||||
"""Test empty() correctly handles None queue."""
|
||||
assert action_queue_rtc_enabled.queue is None
|
||||
assert action_queue_rtc_enabled.empty() is True
|
||||
|
||||
|
||||
def test_qsize_with_none_queue(action_queue_rtc_enabled):
|
||||
"""Test qsize() correctly handles None queue."""
|
||||
assert action_queue_rtc_enabled.queue is None
|
||||
assert action_queue_rtc_enabled.qsize() == 0
|
||||
|
||||
|
||||
# ====================== Integration Tests ======================
|
||||
|
||||
|
||||
def test_typical_rtc_workflow(action_queue_rtc_enabled, sample_actions):
|
||||
"""Test a typical RTC workflow: merge, consume, merge with delay."""
|
||||
# First inference
|
||||
action_queue_rtc_enabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
initial_size = action_queue_rtc_enabled.qsize()
|
||||
assert initial_size == 50
|
||||
|
||||
# Consume 10 actions (execution_horizon)
|
||||
for _ in range(10):
|
||||
action = action_queue_rtc_enabled.get()
|
||||
assert action is not None
|
||||
|
||||
assert action_queue_rtc_enabled.qsize() == 40
|
||||
|
||||
# Second inference with delay
|
||||
action_index_before = action_queue_rtc_enabled.get_action_index()
|
||||
|
||||
action_queue_rtc_enabled.merge(
|
||||
sample_actions["original"],
|
||||
sample_actions["processed"],
|
||||
real_delay=5,
|
||||
action_index_before_inference=action_index_before,
|
||||
)
|
||||
|
||||
# Queue should be replaced, minus delay
|
||||
assert action_queue_rtc_enabled.qsize() == 45
|
||||
assert action_queue_rtc_enabled.get_action_index() == 0
|
||||
|
||||
|
||||
def test_typical_non_rtc_workflow(action_queue_rtc_disabled, sample_actions):
|
||||
"""Test a typical non-RTC workflow: merge, consume, merge again."""
|
||||
# First inference
|
||||
action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
assert action_queue_rtc_disabled.qsize() == 50
|
||||
|
||||
# Consume 40 actions
|
||||
for _ in range(40):
|
||||
action = action_queue_rtc_disabled.get()
|
||||
assert action is not None
|
||||
|
||||
assert action_queue_rtc_disabled.qsize() == 10
|
||||
|
||||
# Second inference (should append)
|
||||
action_queue_rtc_disabled.merge(sample_actions["original"], sample_actions["processed"], real_delay=0)
|
||||
|
||||
# Should have 10 remaining + 50 new = 60
|
||||
assert action_queue_rtc_disabled.qsize() == 60
|
||||
@@ -0,0 +1,65 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for RTC configuration module."""
|
||||
|
||||
from lerobot.configs.types import RTCAttentionSchedule
|
||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
||||
|
||||
# ====================== Initialization Tests ======================
|
||||
|
||||
|
||||
def test_rtc_config_default_initialization():
|
||||
"""Test RTCConfig initializes with default values."""
|
||||
config = RTCConfig()
|
||||
|
||||
assert config.enabled is False
|
||||
assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
|
||||
assert config.max_guidance_weight == 10.0
|
||||
assert config.execution_horizon == 10
|
||||
assert config.debug is False
|
||||
assert config.debug_maxlen == 100
|
||||
|
||||
|
||||
def test_rtc_config_custom_initialization():
|
||||
"""Test RTCConfig initializes with custom values."""
|
||||
config = RTCConfig(
|
||||
enabled=True,
|
||||
prefix_attention_schedule=RTCAttentionSchedule.EXP,
|
||||
max_guidance_weight=5.0,
|
||||
execution_horizon=20,
|
||||
debug=True,
|
||||
debug_maxlen=200,
|
||||
)
|
||||
|
||||
assert config.enabled is True
|
||||
assert config.prefix_attention_schedule == RTCAttentionSchedule.EXP
|
||||
assert config.max_guidance_weight == 5.0
|
||||
assert config.execution_horizon == 20
|
||||
assert config.debug is True
|
||||
assert config.debug_maxlen == 200
|
||||
|
||||
|
||||
def test_rtc_config_partial_initialization():
|
||||
"""Test RTCConfig with partial custom values."""
|
||||
config = RTCConfig(enabled=True, max_guidance_weight=15.0)
|
||||
|
||||
assert config.enabled is True
|
||||
assert config.max_guidance_weight == 15.0
|
||||
# Other values should be defaults
|
||||
assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
|
||||
assert config.execution_horizon == 10
|
||||
assert config.debug is False
|
||||
@@ -0,0 +1,488 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for RTC debug tracker module."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.rtc.debug_tracker import DebugStep, Tracker
|
||||
|
||||
# ====================== Fixtures ======================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_tensors():
|
||||
"""Create sample tensors for testing."""
|
||||
return {
|
||||
"x_t": torch.randn(1, 50, 6),
|
||||
"v_t": torch.randn(1, 50, 6),
|
||||
"x1_t": torch.randn(1, 50, 6),
|
||||
"correction": torch.randn(1, 50, 6),
|
||||
"err": torch.randn(1, 50, 6),
|
||||
"weights": torch.randn(1, 50, 1),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def enabled_tracker():
|
||||
"""Create an enabled tracker with default settings."""
|
||||
return Tracker(enabled=True, maxlen=100)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def disabled_tracker():
|
||||
"""Create a disabled tracker."""
|
||||
return Tracker(enabled=False)
|
||||
|
||||
|
||||
# ====================== DebugStep Tests ======================
|
||||
|
||||
|
||||
def test_debug_step_initialization():
|
||||
"""Test that DebugStep can be initialized with default values."""
|
||||
step = DebugStep()
|
||||
assert step.step_idx == 0
|
||||
assert step.x_t is None
|
||||
assert step.v_t is None
|
||||
assert step.x1_t is None
|
||||
assert step.correction is None
|
||||
assert step.err is None
|
||||
assert step.weights is None
|
||||
assert step.guidance_weight is None
|
||||
assert step.time is None
|
||||
assert step.inference_delay is None
|
||||
assert step.execution_horizon is None
|
||||
assert step.metadata == {}
|
||||
|
||||
|
||||
def test_debug_step_with_values(sample_tensors):
|
||||
"""Test DebugStep initialization with actual values."""
|
||||
step = DebugStep(
|
||||
step_idx=5,
|
||||
x_t=sample_tensors["x_t"],
|
||||
v_t=sample_tensors["v_t"],
|
||||
x1_t=sample_tensors["x1_t"],
|
||||
correction=sample_tensors["correction"],
|
||||
err=sample_tensors["err"],
|
||||
weights=sample_tensors["weights"],
|
||||
guidance_weight=2.5,
|
||||
time=0.8,
|
||||
inference_delay=4,
|
||||
execution_horizon=8,
|
||||
metadata={"custom_key": "custom_value"},
|
||||
)
|
||||
|
||||
assert step.step_idx == 5
|
||||
assert torch.equal(step.x_t, sample_tensors["x_t"])
|
||||
assert torch.equal(step.v_t, sample_tensors["v_t"])
|
||||
assert torch.equal(step.x1_t, sample_tensors["x1_t"])
|
||||
assert torch.equal(step.correction, sample_tensors["correction"])
|
||||
assert torch.equal(step.err, sample_tensors["err"])
|
||||
assert torch.equal(step.weights, sample_tensors["weights"])
|
||||
assert step.guidance_weight == 2.5
|
||||
assert step.time == 0.8
|
||||
assert step.inference_delay == 4
|
||||
assert step.execution_horizon == 8
|
||||
assert step.metadata == {"custom_key": "custom_value"}
|
||||
|
||||
|
||||
def test_debug_step_to_dict_without_tensors(sample_tensors):
|
||||
"""Test converting DebugStep to dictionary without tensor values."""
|
||||
step = DebugStep(
|
||||
step_idx=3,
|
||||
x_t=sample_tensors["x_t"],
|
||||
v_t=sample_tensors["v_t"],
|
||||
guidance_weight=torch.tensor(3.0),
|
||||
time=torch.tensor(0.5),
|
||||
inference_delay=2,
|
||||
execution_horizon=10,
|
||||
)
|
||||
|
||||
result = step.to_dict(include_tensors=False)
|
||||
|
||||
assert result["step_idx"] == 3
|
||||
assert result["guidance_weight"] == 3.0
|
||||
assert result["time"] == 0.5
|
||||
assert result["inference_delay"] == 2
|
||||
assert result["execution_horizon"] == 10
|
||||
|
||||
# Check tensor statistics are included
|
||||
assert "x_t_stats" in result
|
||||
assert "v_t_stats" in result
|
||||
assert "x1_t_stats" not in result # x1_t was None
|
||||
|
||||
# Verify statistics structure
|
||||
assert "shape" in result["x_t_stats"]
|
||||
assert "mean" in result["x_t_stats"]
|
||||
assert "std" in result["x_t_stats"]
|
||||
assert "min" in result["x_t_stats"]
|
||||
assert "max" in result["x_t_stats"]
|
||||
|
||||
# Verify shape matches original tensor
|
||||
assert result["x_t_stats"]["shape"] == tuple(sample_tensors["x_t"].shape)
|
||||
|
||||
|
||||
def test_debug_step_to_dict_with_tensors(sample_tensors):
|
||||
"""Test converting DebugStep to dictionary with tensor values."""
|
||||
step = DebugStep(
|
||||
step_idx=1,
|
||||
x_t=sample_tensors["x_t"],
|
||||
v_t=sample_tensors["v_t"],
|
||||
guidance_weight=1.5,
|
||||
time=0.9,
|
||||
)
|
||||
|
||||
result = step.to_dict(include_tensors=True)
|
||||
|
||||
assert result["step_idx"] == 1
|
||||
assert result["guidance_weight"] == 1.5
|
||||
assert result["time"] == 0.9
|
||||
|
||||
# Check tensors are included (as CPU tensors)
|
||||
assert "x_t" in result
|
||||
assert "v_t" in result
|
||||
assert isinstance(result["x_t"], torch.Tensor)
|
||||
assert isinstance(result["v_t"], torch.Tensor)
|
||||
assert result["x_t"].device.type == "cpu"
|
||||
assert result["v_t"].device.type == "cpu"
|
||||
|
||||
|
||||
def test_debug_step_to_dict_with_none_guidance_weight():
|
||||
"""Test to_dict handles None guidance_weight correctly."""
|
||||
step = DebugStep(step_idx=0, time=1.0, guidance_weight=None)
|
||||
result = step.to_dict(include_tensors=False)
|
||||
assert result["guidance_weight"] is None
|
||||
|
||||
|
||||
def test_tracker_initialization_enabled():
|
||||
"""Test tracker initialization when enabled."""
|
||||
tracker = Tracker(enabled=True, maxlen=50)
|
||||
assert tracker.enabled is True
|
||||
assert tracker._steps == {}
|
||||
assert tracker._maxlen == 50
|
||||
assert tracker._step_counter == 0
|
||||
assert len(tracker) == 0
|
||||
|
||||
|
||||
def test_tracker_reset_when_enabled(enabled_tracker, sample_tensors):
|
||||
"""Test reset clears all steps when tracker is enabled."""
|
||||
# Add some steps
|
||||
enabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
|
||||
enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
|
||||
assert len(enabled_tracker) == 2
|
||||
|
||||
# Reset
|
||||
enabled_tracker.reset()
|
||||
assert len(enabled_tracker) == 0
|
||||
assert enabled_tracker._step_counter == 0
|
||||
assert enabled_tracker._steps == {}
|
||||
|
||||
|
||||
def test_tracker_reset_when_disabled(disabled_tracker):
|
||||
"""Test reset on disabled tracker doesn't cause errors."""
|
||||
disabled_tracker.reset()
|
||||
assert len(disabled_tracker) == 0
|
||||
|
||||
|
||||
# ====================== Tracker.track() Tests ======================
|
||||
|
||||
|
||||
def test_track_creates_new_step(enabled_tracker, sample_tensors):
|
||||
"""Test that track creates a new step when time doesn't exist."""
|
||||
enabled_tracker.track(
|
||||
time=1.0,
|
||||
x_t=sample_tensors["x_t"],
|
||||
v_t=sample_tensors["v_t"],
|
||||
guidance_weight=5.0,
|
||||
inference_delay=4,
|
||||
execution_horizon=8,
|
||||
)
|
||||
|
||||
assert len(enabled_tracker) == 1
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert len(steps) == 1
|
||||
assert steps[0].step_idx == 0
|
||||
assert steps[0].time == 1.0
|
||||
assert torch.equal(steps[0].x_t, sample_tensors["x_t"])
|
||||
assert torch.equal(steps[0].v_t, sample_tensors["v_t"])
|
||||
assert steps[0].guidance_weight == 5.0
|
||||
assert steps[0].inference_delay == 4
|
||||
assert steps[0].execution_horizon == 8
|
||||
|
||||
|
||||
def test_track_updates_existing_step(enabled_tracker, sample_tensors):
|
||||
"""Test that track updates an existing step at the same time."""
|
||||
# Create initial step
|
||||
enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
|
||||
assert len(enabled_tracker) == 1
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert steps[0].v_t is None
|
||||
|
||||
# Update the same timestep with v_t
|
||||
enabled_tracker.track(time=0.9, v_t=sample_tensors["v_t"])
|
||||
assert len(enabled_tracker) == 1 # Still only one step
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert torch.equal(steps[0].x_t, sample_tensors["x_t"]) # Original x_t preserved
|
||||
assert torch.equal(steps[0].v_t, sample_tensors["v_t"]) # New v_t added
|
||||
|
||||
|
||||
def test_track_with_tensor_time(enabled_tracker, sample_tensors):
|
||||
"""Test track handles tensor time values correctly."""
|
||||
time_tensor = torch.tensor(0.8)
|
||||
enabled_tracker.track(time=time_tensor, x_t=sample_tensors["x_t"])
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert len(steps) == 1
|
||||
assert abs(steps[0].time - 0.8) < 1e-6 # Use approximate comparison for floating point
|
||||
|
||||
|
||||
def test_track_time_rounding(enabled_tracker, sample_tensors):
|
||||
"""Test that track rounds time to avoid floating point precision issues."""
|
||||
# These times should be treated as the same after rounding to 6 decimals
|
||||
enabled_tracker.track(time=0.9000001, x_t=sample_tensors["x_t"])
|
||||
enabled_tracker.track(time=0.9000002, v_t=sample_tensors["v_t"])
|
||||
|
||||
# Should still be one step (times rounded to same value)
|
||||
assert len(enabled_tracker) == 1
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert torch.equal(steps[0].x_t, sample_tensors["x_t"])
|
||||
assert torch.equal(steps[0].v_t, sample_tensors["v_t"])
|
||||
|
||||
|
||||
def test_track_does_nothing_when_disabled(disabled_tracker, sample_tensors):
|
||||
"""Test that track does nothing when tracker is disabled."""
|
||||
disabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
|
||||
assert len(disabled_tracker) == 0
|
||||
|
||||
|
||||
def test_track_with_metadata(enabled_tracker, sample_tensors):
|
||||
"""Test track stores custom metadata."""
|
||||
enabled_tracker.track(time=0.7, x_t=sample_tensors["x_t"], custom_field="custom_value", count=42)
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert steps[0].metadata["custom_field"] == "custom_value"
|
||||
assert steps[0].metadata["count"] == 42
|
||||
|
||||
|
||||
def test_track_updates_metadata(enabled_tracker):
|
||||
"""Test that track updates metadata for existing steps."""
|
||||
enabled_tracker.track(time=0.6, meta1="value1")
|
||||
enabled_tracker.track(time=0.6, meta2="value2")
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert steps[0].metadata["meta1"] == "value1"
|
||||
assert steps[0].metadata["meta2"] == "value2"
|
||||
|
||||
|
||||
def test_track_clones_tensors(enabled_tracker, sample_tensors):
|
||||
"""Test that track clones tensors instead of storing references."""
|
||||
x_t_original = sample_tensors["x_t"].clone()
|
||||
enabled_tracker.track(time=0.5, x_t=sample_tensors["x_t"])
|
||||
|
||||
# Modify original tensor
|
||||
sample_tensors["x_t"].fill_(999.0)
|
||||
|
||||
# Tracked tensor should not be affected
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert not torch.equal(steps[0].x_t, sample_tensors["x_t"])
|
||||
assert torch.equal(steps[0].x_t, x_t_original)
|
||||
|
||||
|
||||
def test_track_with_none_values(enabled_tracker):
|
||||
"""Test track handles None values correctly."""
|
||||
enabled_tracker.track(
|
||||
time=0.4,
|
||||
x_t=None,
|
||||
v_t=None,
|
||||
guidance_weight=None,
|
||||
inference_delay=None,
|
||||
)
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert len(steps) == 1
|
||||
assert steps[0].x_t is None
|
||||
assert steps[0].v_t is None
|
||||
assert steps[0].guidance_weight is None
|
||||
assert steps[0].inference_delay is None
|
||||
|
||||
|
||||
def test_track_updates_only_non_none_fields(enabled_tracker, sample_tensors):
|
||||
"""Test that update preserves existing values when None is passed."""
|
||||
# Create step with x_t
|
||||
enabled_tracker.track(time=0.3, x_t=sample_tensors["x_t"], guidance_weight=2.0)
|
||||
|
||||
# Update with v_t only (pass None for other fields)
|
||||
enabled_tracker.track(time=0.3, v_t=sample_tensors["v_t"], x_t=None, guidance_weight=None)
|
||||
|
||||
# Original values should be preserved
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert torch.equal(steps[0].x_t, sample_tensors["x_t"]) # Still has x_t
|
||||
assert torch.equal(steps[0].v_t, sample_tensors["v_t"]) # Now has v_t
|
||||
assert steps[0].guidance_weight == 2.0 # Still has guidance_weight
|
||||
|
||||
|
||||
# ====================== Tracker.maxlen Tests ======================
|
||||
|
||||
|
||||
def test_tracker_enforces_maxlen():
|
||||
"""Test that tracker enforces maxlen limit."""
|
||||
tracker = Tracker(enabled=True, maxlen=3)
|
||||
|
||||
# Add 5 steps
|
||||
for i in range(5):
|
||||
time = 1.0 - i * 0.1 # 1.0, 0.9, 0.8, 0.7, 0.6
|
||||
tracker.track(time=time, x_t=torch.randn(1, 10, 6))
|
||||
|
||||
# Should only keep the last 3
|
||||
assert len(tracker) == 3
|
||||
|
||||
# Verify oldest steps were removed (should have 0.6, 0.7, 0.8)
|
||||
steps = tracker.get_all_steps()
|
||||
times = sorted([step.time for step in steps])
|
||||
assert times == [0.6, 0.7, 0.8]
|
||||
|
||||
|
||||
def test_tracker_step_idx_increments_despite_maxlen():
|
||||
"""Test that step_idx continues incrementing even when maxlen is enforced."""
|
||||
tracker = Tracker(enabled=True, maxlen=2)
|
||||
|
||||
# Add 4 steps
|
||||
for i in range(4):
|
||||
time = 1.0 - i * 0.1
|
||||
tracker.track(time=time, x_t=torch.randn(1, 10, 6))
|
||||
|
||||
# Should have 2 steps with step_idx 2 and 3 (oldest removed)
|
||||
steps = sorted(tracker.get_all_steps(), key=lambda s: s.step_idx)
|
||||
assert len(steps) == 2
|
||||
assert steps[0].step_idx == 2
|
||||
assert steps[1].step_idx == 3
|
||||
|
||||
|
||||
def test_tracker_without_maxlen_keeps_all():
|
||||
"""Test that tracker without maxlen keeps all steps."""
|
||||
tracker = Tracker(enabled=True, maxlen=None)
|
||||
|
||||
# Add 100 steps
|
||||
for i in range(100):
|
||||
time = 1.0 - i * 0.01
|
||||
tracker.track(time=time, x_t=torch.randn(1, 10, 6))
|
||||
|
||||
assert len(tracker) == 100
|
||||
|
||||
|
||||
def test_get_all_steps_returns_empty_when_disabled(disabled_tracker):
|
||||
"""Test get_all_steps returns empty list when disabled."""
|
||||
steps = disabled_tracker.get_all_steps()
|
||||
assert steps == []
|
||||
assert isinstance(steps, list)
|
||||
|
||||
|
||||
def test_get_all_steps_returns_empty_when_no_steps(enabled_tracker):
|
||||
"""Test get_all_steps returns empty list when no steps tracked."""
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert steps == []
|
||||
|
||||
|
||||
def test_get_all_steps_returns_all_tracked_steps(enabled_tracker, sample_tensors):
|
||||
"""Test get_all_steps returns all tracked steps."""
|
||||
# Track 5 steps
|
||||
for i in range(5):
|
||||
time = 1.0 - i * 0.1
|
||||
enabled_tracker.track(time=time, x_t=sample_tensors["x_t"])
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert len(steps) == 5
|
||||
|
||||
# Verify all are DebugStep instances
|
||||
for step in steps:
|
||||
assert isinstance(step, DebugStep)
|
||||
|
||||
|
||||
def test_get_all_steps_preserves_insertion_order(enabled_tracker):
|
||||
"""Test that get_all_steps preserves insertion order (Python 3.7+)."""
|
||||
times = [0.9, 0.8, 0.7, 0.6, 0.5]
|
||||
for time in times:
|
||||
enabled_tracker.track(time=time, x_t=torch.randn(1, 10, 6))
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
retrieved_times = [step.time for step in steps]
|
||||
|
||||
# Should be in insertion order
|
||||
assert retrieved_times == times
|
||||
|
||||
|
||||
# ====================== Tracker.__len__() Tests ======================
|
||||
|
||||
|
||||
def test_len_returns_zero_when_disabled(disabled_tracker):
|
||||
"""Test __len__ returns 0 when tracker is disabled."""
|
||||
assert len(disabled_tracker) == 0
|
||||
|
||||
|
||||
def test_len_returns_zero_when_empty(enabled_tracker):
|
||||
"""Test __len__ returns 0 when no steps are tracked."""
|
||||
assert len(enabled_tracker) == 0
|
||||
|
||||
|
||||
def test_len_returns_correct_count(enabled_tracker, sample_tensors):
|
||||
"""Test __len__ returns correct number of tracked steps."""
|
||||
assert len(enabled_tracker) == 0
|
||||
|
||||
enabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
|
||||
assert len(enabled_tracker) == 1
|
||||
|
||||
enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
|
||||
assert len(enabled_tracker) == 2
|
||||
|
||||
enabled_tracker.track(time=0.8, x_t=sample_tensors["x_t"])
|
||||
assert len(enabled_tracker) == 3
|
||||
|
||||
|
||||
def test_len_after_reset(enabled_tracker, sample_tensors):
|
||||
"""Test __len__ returns 0 after reset."""
|
||||
enabled_tracker.track(time=1.0, x_t=sample_tensors["x_t"])
|
||||
enabled_tracker.track(time=0.9, x_t=sample_tensors["x_t"])
|
||||
assert len(enabled_tracker) == 2
|
||||
|
||||
enabled_tracker.reset()
|
||||
assert len(enabled_tracker) == 0
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_tracker_handles_gpu_tensors():
|
||||
"""Test tracker correctly handles GPU tensors."""
|
||||
tracker = Tracker(enabled=True, maxlen=10)
|
||||
x_t_gpu = torch.randn(1, 50, 6, device="cuda")
|
||||
|
||||
tracker.track(time=1.0, x_t=x_t_gpu)
|
||||
|
||||
steps = tracker.get_all_steps()
|
||||
# Tracker should clone and detach tensors
|
||||
assert steps[0].x_t.device.type == "cuda"
|
||||
|
||||
|
||||
def test_tracker_with_varying_tensor_shapes(enabled_tracker):
|
||||
"""Test tracker handles varying tensor shapes across steps."""
|
||||
enabled_tracker.track(time=1.0, x_t=torch.randn(1, 50, 6))
|
||||
enabled_tracker.track(time=0.9, x_t=torch.randn(1, 25, 6))
|
||||
enabled_tracker.track(time=0.8, x_t=torch.randn(2, 50, 8))
|
||||
|
||||
steps = enabled_tracker.get_all_steps()
|
||||
assert len(steps) == 3
|
||||
assert steps[0].x_t.shape == (1, 50, 6)
|
||||
assert steps[1].x_t.shape == (1, 25, 6)
|
||||
assert steps[2].x_t.shape == (2, 50, 8)
|
||||
@@ -0,0 +1,322 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for RTC LatencyTracker module."""
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.policies.rtc.latency_tracker import LatencyTracker
|
||||
|
||||
# ====================== Fixtures ======================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tracker():
|
||||
"""Create a LatencyTracker with default maxlen."""
|
||||
return LatencyTracker(maxlen=100)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def small_tracker():
|
||||
"""Create a LatencyTracker with small maxlen for overflow testing."""
|
||||
return LatencyTracker(maxlen=5)
|
||||
|
||||
|
||||
# ====================== Initialization Tests ======================
|
||||
|
||||
|
||||
def test_latency_tracker_initialization():
|
||||
"""Test LatencyTracker initializes correctly."""
|
||||
tracker = LatencyTracker(maxlen=50)
|
||||
assert len(tracker) == 0
|
||||
assert tracker.max_latency == 0.0
|
||||
assert tracker.max() == 0.0
|
||||
|
||||
|
||||
def test_latency_tracker_default_maxlen():
|
||||
"""Test LatencyTracker uses default maxlen."""
|
||||
tracker = LatencyTracker()
|
||||
# Should accept default maxlen=100
|
||||
assert len(tracker) == 0
|
||||
|
||||
|
||||
# ====================== add() Tests ======================
|
||||
|
||||
|
||||
def test_add_single_latency(tracker):
|
||||
"""Test adding a single latency value."""
|
||||
tracker.add(0.5)
|
||||
assert len(tracker) == 1
|
||||
assert tracker.max() == 0.5
|
||||
|
||||
|
||||
def test_add_multiple_latencies(tracker):
|
||||
"""Test adding multiple latency values."""
|
||||
latencies = [0.1, 0.5, 0.3, 0.8, 0.2]
|
||||
for lat in latencies:
|
||||
tracker.add(lat)
|
||||
|
||||
assert len(tracker) == 5
|
||||
assert tracker.max() == 0.8
|
||||
|
||||
|
||||
def test_add_negative_latency_ignored(tracker):
|
||||
"""Test that negative latencies are ignored."""
|
||||
tracker.add(0.5)
|
||||
tracker.add(-0.1)
|
||||
tracker.add(0.3)
|
||||
|
||||
# Should only have 2 valid latencies
|
||||
assert len(tracker) == 2
|
||||
assert tracker.max() == 0.5
|
||||
|
||||
|
||||
def test_add_zero_latency(tracker):
|
||||
"""Test adding zero latency."""
|
||||
tracker.add(0.0)
|
||||
assert len(tracker) == 1
|
||||
assert tracker.max() == 0.0
|
||||
|
||||
|
||||
def test_add_converts_to_float(tracker):
|
||||
"""Test add() converts input to float."""
|
||||
tracker.add(5) # Integer
|
||||
tracker.add("3.5") # String
|
||||
|
||||
assert len(tracker) == 2
|
||||
assert tracker.max() == 5.0
|
||||
|
||||
|
||||
def test_add_updates_max_latency(tracker):
|
||||
"""Test that max_latency is updated correctly."""
|
||||
tracker.add(0.5)
|
||||
assert tracker.max_latency == 0.5
|
||||
|
||||
tracker.add(0.3)
|
||||
assert tracker.max_latency == 0.5 # Should not decrease
|
||||
|
||||
tracker.add(0.9)
|
||||
assert tracker.max_latency == 0.9 # Should increase
|
||||
|
||||
|
||||
# ====================== reset() Tests ======================
|
||||
|
||||
|
||||
def test_reset_clears_values(tracker):
|
||||
"""Test reset() clears all values."""
|
||||
tracker.add(0.5)
|
||||
tracker.add(0.8)
|
||||
tracker.add(0.3)
|
||||
assert len(tracker) == 3
|
||||
|
||||
tracker.reset()
|
||||
assert len(tracker) == 0
|
||||
assert tracker.max_latency == 0.0
|
||||
|
||||
|
||||
def test_reset_clears_max_latency(tracker):
|
||||
"""Test reset() resets max_latency."""
|
||||
tracker.add(1.5)
|
||||
assert tracker.max_latency == 1.5
|
||||
|
||||
tracker.reset()
|
||||
assert tracker.max_latency == 0.0
|
||||
|
||||
|
||||
def test_reset_allows_new_values(tracker):
|
||||
"""Test that tracker works correctly after reset."""
|
||||
tracker.add(0.5)
|
||||
tracker.reset()
|
||||
|
||||
tracker.add(0.3)
|
||||
assert len(tracker) == 1
|
||||
assert tracker.max() == 0.3
|
||||
|
||||
|
||||
# ====================== max() Tests ======================
|
||||
|
||||
|
||||
def test_max_returns_zero_when_empty(tracker):
|
||||
"""Test max() returns 0.0 when tracker is empty."""
|
||||
assert tracker.max() == 0.0
|
||||
|
||||
|
||||
def test_max_returns_maximum_value(tracker):
|
||||
"""Test max() returns the maximum latency."""
|
||||
latencies = [0.2, 0.8, 0.3, 0.5, 0.1]
|
||||
for lat in latencies:
|
||||
tracker.add(lat)
|
||||
|
||||
assert tracker.max() == 0.8
|
||||
|
||||
|
||||
def test_max_persists_after_sliding_window(small_tracker):
|
||||
"""Test max() persists even after values slide out of window."""
|
||||
# Add values that will exceed maxlen=5
|
||||
small_tracker.add(0.1)
|
||||
small_tracker.add(0.9) # This is max
|
||||
small_tracker.add(0.2)
|
||||
small_tracker.add(0.3)
|
||||
small_tracker.add(0.4)
|
||||
small_tracker.add(0.5) # This pushes out 0.1
|
||||
|
||||
# Max should still be 0.9 even though only last 5 values kept
|
||||
assert small_tracker.max() == 0.9
|
||||
|
||||
|
||||
def test_max_after_reset(tracker):
|
||||
"""Test max() returns 0.0 after reset."""
|
||||
tracker.add(1.5)
|
||||
tracker.reset()
|
||||
assert tracker.max() == 0.0
|
||||
|
||||
|
||||
# ====================== p95() Tests ======================
|
||||
|
||||
|
||||
def test_p95_returns_zero_when_empty(tracker):
|
||||
"""Test p95() returns 0.0 when tracker is empty."""
|
||||
assert tracker.p95() == 0.0
|
||||
|
||||
|
||||
def test_p95_returns_95th_percentile(tracker):
|
||||
"""Test p95() returns the 95th percentile."""
|
||||
# Add 100 values
|
||||
for i in range(100):
|
||||
tracker.add(i / 100.0)
|
||||
|
||||
p95 = tracker.p95()
|
||||
assert 0.93 <= p95 <= 0.96
|
||||
|
||||
|
||||
def test_p95_equals_percentile_95(tracker):
|
||||
"""Test p95() equals percentile(0.95)."""
|
||||
for i in range(50):
|
||||
tracker.add(i / 50.0)
|
||||
|
||||
assert tracker.p95() == tracker.percentile(0.95)
|
||||
|
||||
|
||||
# ====================== Edge Cases Tests ======================
|
||||
|
||||
|
||||
def test_single_value(tracker):
|
||||
"""Test tracker behavior with single value."""
|
||||
tracker.add(0.75)
|
||||
|
||||
assert len(tracker) == 1
|
||||
assert tracker.max() == 0.75
|
||||
assert tracker.percentile(0.0) == 0.75
|
||||
assert tracker.percentile(0.5) == 0.75
|
||||
assert tracker.percentile(1.0) == 0.75
|
||||
|
||||
|
||||
def test_all_same_values(tracker):
|
||||
"""Test tracker with all identical values."""
|
||||
for _ in range(10):
|
||||
tracker.add(0.5)
|
||||
|
||||
assert len(tracker) == 10
|
||||
assert tracker.max() == 0.5
|
||||
assert tracker.percentile(0.0) == 0.5
|
||||
assert tracker.percentile(0.5) == 0.5
|
||||
assert tracker.percentile(1.0) == 0.5
|
||||
|
||||
|
||||
def test_very_small_values(tracker):
|
||||
"""Test tracker with very small float values."""
|
||||
tracker.add(1e-10)
|
||||
tracker.add(2e-10)
|
||||
tracker.add(3e-10)
|
||||
|
||||
assert len(tracker) == 3
|
||||
assert tracker.max() == pytest.approx(3e-10)
|
||||
|
||||
|
||||
def test_very_large_values(tracker):
|
||||
"""Test tracker with very large float values."""
|
||||
tracker.add(1e10)
|
||||
tracker.add(2e10)
|
||||
tracker.add(3e10)
|
||||
|
||||
assert len(tracker) == 3
|
||||
assert tracker.max() == pytest.approx(3e10)
|
||||
|
||||
|
||||
# ====================== Integration Tests ======================
|
||||
|
||||
|
||||
def test_typical_usage_pattern(tracker):
|
||||
"""Test a typical usage pattern of the tracker."""
|
||||
# Simulate adding latencies over time
|
||||
latencies = [0.05, 0.08, 0.12, 0.07, 0.15, 0.09, 0.11, 0.06, 0.14, 0.10]
|
||||
|
||||
for lat in latencies:
|
||||
tracker.add(lat)
|
||||
|
||||
# Check statistics
|
||||
assert len(tracker) == 10
|
||||
assert tracker.max() == 0.15
|
||||
|
||||
# p95 should be close to max since we have only 10 values
|
||||
p95 = tracker.p95()
|
||||
assert p95 >= tracker.percentile(0.5) # p95 should be >= median
|
||||
assert p95 <= tracker.max() # p95 should be <= max
|
||||
|
||||
|
||||
def test_reset_and_reuse(tracker):
|
||||
"""Test resetting and reusing tracker."""
|
||||
# First batch
|
||||
tracker.add(1.0)
|
||||
tracker.add(2.0)
|
||||
assert tracker.max() == 2.0
|
||||
|
||||
# Reset
|
||||
tracker.reset()
|
||||
|
||||
# Second batch
|
||||
tracker.add(0.5)
|
||||
tracker.add(0.8)
|
||||
assert len(tracker) == 2
|
||||
assert tracker.max() == 0.8
|
||||
assert tracker.percentile(0.5) <= 0.8
|
||||
|
||||
|
||||
# ====================== Type Conversion Tests ======================
|
||||
|
||||
|
||||
def test_add_with_integer(tracker):
|
||||
"""Test adding integer values."""
|
||||
tracker.add(5)
|
||||
assert len(tracker) == 1
|
||||
assert tracker.max() == 5.0
|
||||
|
||||
|
||||
def test_add_with_string_number(tracker):
|
||||
"""Test adding string representation of number."""
|
||||
tracker.add("3.14")
|
||||
assert len(tracker) == 1
|
||||
assert tracker.max() == pytest.approx(3.14)
|
||||
|
||||
|
||||
def test_percentile_converts_q_to_float(tracker):
|
||||
"""Test percentile converts q parameter to float."""
|
||||
tracker.add(0.5)
|
||||
tracker.add(0.8)
|
||||
|
||||
# Pass integer q
|
||||
result = tracker.percentile(1)
|
||||
assert result == 0.8
|
||||
@@ -0,0 +1,773 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for RTC modeling module (RTCProcessor)."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import RTCAttentionSchedule
|
||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
||||
from lerobot.policies.rtc.modeling_rtc import RTCProcessor
|
||||
|
||||
# ====================== Fixtures ======================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtc_config_debug_enabled():
|
||||
"""Create RTC config with debug enabled."""
|
||||
return RTCConfig(
|
||||
enabled=True,
|
||||
prefix_attention_schedule=RTCAttentionSchedule.LINEAR,
|
||||
max_guidance_weight=10.0,
|
||||
execution_horizon=10,
|
||||
debug=True,
|
||||
debug_maxlen=100,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtc_config_debug_disabled():
|
||||
"""Create RTC config with debug disabled."""
|
||||
return RTCConfig(
|
||||
enabled=True,
|
||||
prefix_attention_schedule=RTCAttentionSchedule.LINEAR,
|
||||
max_guidance_weight=10.0,
|
||||
execution_horizon=10,
|
||||
debug=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtc_processor_debug_enabled(rtc_config_debug_enabled):
|
||||
"""Create RTCProcessor with debug enabled."""
|
||||
return RTCProcessor(rtc_config_debug_enabled)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rtc_processor_debug_disabled(rtc_config_debug_disabled):
|
||||
"""Create RTCProcessor with debug disabled."""
|
||||
return RTCProcessor(rtc_config_debug_disabled)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_x_t():
|
||||
"""Create sample x_t tensor (batch, time, action_dim)."""
|
||||
return torch.randn(1, 50, 6)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_prev_chunk():
|
||||
"""Create sample previous chunk tensor."""
|
||||
return torch.randn(1, 50, 6)
|
||||
|
||||
|
||||
# ====================== Initialization Tests ======================
|
||||
|
||||
|
||||
def test_rtc_processor_initialization_with_debug(rtc_config_debug_enabled):
|
||||
"""Test RTCProcessor initializes with debug tracker."""
|
||||
processor = RTCProcessor(rtc_config_debug_enabled)
|
||||
assert processor.rtc_config == rtc_config_debug_enabled
|
||||
assert processor.tracker is not None
|
||||
assert processor.tracker.enabled is True
|
||||
|
||||
|
||||
def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled):
|
||||
"""Test RTCProcessor initializes without debug tracker."""
|
||||
processor = RTCProcessor(rtc_config_debug_disabled)
|
||||
assert processor.rtc_config == rtc_config_debug_disabled
|
||||
assert processor.tracker is None
|
||||
|
||||
|
||||
# ====================== Tracker Proxy Methods Tests ======================
|
||||
|
||||
|
||||
def test_track_when_tracker_enabled(rtc_processor_debug_enabled, sample_x_t):
|
||||
"""Test track() forwards to tracker when enabled."""
|
||||
rtc_processor_debug_enabled.track(
|
||||
time=torch.tensor(0.5),
|
||||
x_t=sample_x_t,
|
||||
v_t=sample_x_t,
|
||||
guidance_weight=2.0,
|
||||
)
|
||||
|
||||
# Should have tracked one step
|
||||
steps = rtc_processor_debug_enabled.get_all_debug_steps()
|
||||
assert len(steps) == 1
|
||||
assert steps[0].time == 0.5
|
||||
|
||||
|
||||
def test_track_when_tracker_disabled(rtc_processor_debug_disabled, sample_x_t):
|
||||
"""Test track() does nothing when tracker disabled."""
|
||||
# Should not raise error
|
||||
rtc_processor_debug_disabled.track(
|
||||
time=torch.tensor(0.5),
|
||||
x_t=sample_x_t,
|
||||
v_t=sample_x_t,
|
||||
)
|
||||
|
||||
# Should return empty list
|
||||
steps = rtc_processor_debug_disabled.get_all_debug_steps()
|
||||
assert len(steps) == 0
|
||||
|
||||
|
||||
def test_get_all_debug_steps_when_enabled(rtc_processor_debug_enabled, sample_x_t):
|
||||
"""Test get_all_debug_steps() returns tracked steps."""
|
||||
rtc_processor_debug_enabled.track(time=torch.tensor(0.5), x_t=sample_x_t)
|
||||
rtc_processor_debug_enabled.track(time=torch.tensor(0.4), x_t=sample_x_t)
|
||||
|
||||
steps = rtc_processor_debug_enabled.get_all_debug_steps()
|
||||
assert len(steps) == 2
|
||||
|
||||
|
||||
def test_get_all_debug_steps_when_disabled(rtc_processor_debug_disabled):
|
||||
"""Test get_all_debug_steps() returns empty list when disabled."""
|
||||
steps = rtc_processor_debug_disabled.get_all_debug_steps()
|
||||
assert steps == []
|
||||
assert isinstance(steps, list)
|
||||
|
||||
|
||||
def test_is_debug_enabled_when_tracker_exists(rtc_processor_debug_enabled):
|
||||
"""Test is_debug_enabled() returns True when tracker enabled."""
|
||||
assert rtc_processor_debug_enabled.is_debug_enabled() is True
|
||||
|
||||
|
||||
def test_is_debug_enabled_when_tracker_disabled(rtc_processor_debug_disabled):
|
||||
"""Test is_debug_enabled() returns False when tracker disabled."""
|
||||
assert rtc_processor_debug_disabled.is_debug_enabled() is False
|
||||
|
||||
|
||||
def test_reset_tracker_when_enabled(rtc_processor_debug_enabled, sample_x_t):
|
||||
"""Test reset_tracker() clears tracked steps."""
|
||||
rtc_processor_debug_enabled.track(time=torch.tensor(0.5), x_t=sample_x_t)
|
||||
rtc_processor_debug_enabled.track(time=torch.tensor(0.4), x_t=sample_x_t)
|
||||
assert len(rtc_processor_debug_enabled.get_all_debug_steps()) == 2
|
||||
|
||||
rtc_processor_debug_enabled.reset_tracker()
|
||||
assert len(rtc_processor_debug_enabled.get_all_debug_steps()) == 0
|
||||
|
||||
|
||||
def test_reset_tracker_when_disabled(rtc_processor_debug_disabled):
|
||||
"""Test reset_tracker() doesn't error when tracker disabled."""
|
||||
rtc_processor_debug_disabled.reset_tracker() # Should not raise
|
||||
|
||||
|
||||
# ====================== get_prefix_weights Tests ======================
|
||||
|
||||
|
||||
def test_get_prefix_weights_zeros_schedule():
|
||||
"""Test get_prefix_weights with ZEROS schedule."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.ZEROS)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=5, end=10, total=20)
|
||||
|
||||
# First 5 should be 1.0, rest should be 0.0
|
||||
assert weights.shape == (20,)
|
||||
assert torch.all(weights[:5] == 1.0)
|
||||
assert torch.all(weights[5:] == 0.0)
|
||||
|
||||
|
||||
def test_get_prefix_weights_ones_schedule():
|
||||
"""Test get_prefix_weights with ONES schedule."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.ONES)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=5, end=15, total=20)
|
||||
|
||||
# First 15 should be 1.0, rest should be 0.0
|
||||
assert weights.shape == (20,)
|
||||
assert torch.all(weights[:15] == 1.0)
|
||||
assert torch.all(weights[15:] == 0.0)
|
||||
|
||||
|
||||
def test_get_prefix_weights_linear_schedule():
|
||||
"""Test get_prefix_weights with LINEAR schedule."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=5, end=14, total=25)
|
||||
|
||||
# Should have shape (20,)
|
||||
assert weights.shape == (25,)
|
||||
|
||||
# First 5 should be 1.0 (leading ones)
|
||||
assert torch.all(weights[:5] == 1.0)
|
||||
|
||||
# Middle section (5:15) should be linearly decreasing from 1 to 0
|
||||
middle_weights = torch.tensor([0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])
|
||||
assert torch.allclose(weights[5:14], middle_weights)
|
||||
|
||||
# Last 5 should be 0.0 (trailing zeros)
|
||||
assert torch.all(weights[14:] == 0.0)
|
||||
|
||||
|
||||
def test_get_prefix_weights_exp_schedule():
|
||||
"""Test get_prefix_weights with EXP schedule."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.EXP)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=5, end=14, total=25)
|
||||
|
||||
# Should have shape (20,)
|
||||
assert weights.shape == (25,)
|
||||
|
||||
# First 5 should be 1.0 (leading ones)
|
||||
assert torch.all(weights[:5] == 1.0)
|
||||
|
||||
# Middle section should be exponentially weighted
|
||||
middle_weights = torch.tensor([0.7645, 0.5706, 0.4130, 0.2871, 0.1888, 0.1145, 0.0611, 0.0258, 0.0061])
|
||||
assert torch.allclose(weights[5:14], middle_weights, atol=1e-4)
|
||||
|
||||
# Last 5 should be 0.0 (trailing zeros)
|
||||
assert torch.all(weights[14:] == 0.0)
|
||||
|
||||
|
||||
def test_get_prefix_weights_with_start_equals_end():
|
||||
"""Test get_prefix_weights when start equals end."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=10, end=10, total=20)
|
||||
|
||||
# Should have ones up to start, then zeros
|
||||
assert torch.all(weights[:10] == 1.0)
|
||||
assert torch.all(weights[10:] == 0.0)
|
||||
|
||||
|
||||
def test_get_prefix_weights_with_start_greater_than_end():
|
||||
"""Test get_prefix_weights when start > end (gets clamped)."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
# start > end should use min(start, end) = end
|
||||
weights = processor.get_prefix_weights(start=15, end=10, total=20)
|
||||
|
||||
# Should have ones up to end (10), then zeros
|
||||
assert torch.all(weights[:10] == 1.0)
|
||||
assert torch.all(weights[10:] == 0.0)
|
||||
|
||||
|
||||
# ====================== Helper Method Tests ======================
|
||||
|
||||
|
||||
def test_linweights_with_end_equals_start():
|
||||
"""Test _linweights when end equals start."""
|
||||
config = RTCConfig()
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor._linweights(start=10, end=10, total=20)
|
||||
|
||||
# Should return empty tensor
|
||||
assert len(weights) == 0
|
||||
|
||||
|
||||
def test_linweights_with_end_less_than_start():
|
||||
"""Test _linweights when end < start."""
|
||||
config = RTCConfig()
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor._linweights(start=15, end=10, total=20)
|
||||
|
||||
# Should return empty tensor
|
||||
assert len(weights) == 0
|
||||
|
||||
|
||||
def test_add_trailing_zeros_normal():
|
||||
"""Test _add_trailing_zeros adds zeros correctly."""
|
||||
config = RTCConfig()
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2])
|
||||
result = processor._add_trailing_zeros(weights, total=10, end=5)
|
||||
|
||||
# Should add 5 zeros (total - end = 10 - 5 = 5)
|
||||
assert len(result) == 10
|
||||
assert torch.all(result[:5] == weights)
|
||||
assert torch.all(result[5:] == 0.0)
|
||||
|
||||
|
||||
def test_add_trailing_zeros_no_zeros_needed():
|
||||
"""Test _add_trailing_zeros when no zeros needed."""
|
||||
config = RTCConfig()
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = torch.tensor([1.0, 0.8, 0.6])
|
||||
result = processor._add_trailing_zeros(weights, total=3, end=5)
|
||||
|
||||
# zeros_len = 3 - 5 = -2 <= 0, so no zeros added
|
||||
assert torch.equal(result, weights)
|
||||
|
||||
|
||||
def test_add_leading_ones_normal():
|
||||
"""Test _add_leading_ones adds ones correctly."""
|
||||
config = RTCConfig()
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = torch.tensor([0.8, 0.6, 0.4, 0.2, 0.0])
|
||||
result = processor._add_leading_ones(weights, start=3, total=10)
|
||||
|
||||
# Should add 3 ones at the start
|
||||
assert len(result) == 8
|
||||
assert torch.all(result[:3] == 1.0)
|
||||
assert torch.all(result[3:] == weights)
|
||||
|
||||
|
||||
def test_add_leading_ones_no_ones_needed():
|
||||
"""Test _add_leading_ones when no ones needed."""
|
||||
config = RTCConfig()
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = torch.tensor([0.8, 0.6, 0.4])
|
||||
result = processor._add_leading_ones(weights, start=0, total=10)
|
||||
|
||||
# ones_len = 0, so no ones added
|
||||
assert torch.equal(result, weights)
|
||||
|
||||
|
||||
def test_get_prefix_weights_with_start_equals_total():
|
||||
"""Test get_prefix_weights when start equals total."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=10, end=10, total=20)
|
||||
|
||||
# Should have ones up to start, then zeros
|
||||
assert len(weights) == 20
|
||||
assert torch.all(weights[:10] == 1.0)
|
||||
assert torch.all(weights[10:] == 0.0)
|
||||
|
||||
|
||||
def test_get_prefix_weights_with_total_less_than_start():
|
||||
"""Test get_prefix_weights when total less than start."""
|
||||
config = RTCConfig(prefix_attention_schedule=RTCAttentionSchedule.LINEAR)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
weights = processor.get_prefix_weights(start=10, end=10, total=5)
|
||||
|
||||
# Should have ones up to start, then zeros
|
||||
assert len(weights) == 5
|
||||
assert torch.all(weights == 1.0)
|
||||
|
||||
|
||||
# ====================== denoise_step Tests ======================
|
||||
|
||||
|
||||
def test_denoise_step_without_prev_chunk(rtc_processor_debug_disabled):
|
||||
"""Test denoise_step without previous chunk (no guidance)."""
|
||||
x_t = torch.randn(1, 50, 6)
|
||||
|
||||
# Mock denoiser that returns fixed velocity
|
||||
def mock_denoiser(x):
|
||||
return torch.ones_like(x) * 0.5
|
||||
|
||||
result = rtc_processor_debug_disabled.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=None,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Should return v_t unchanged (no guidance)
|
||||
expected = mock_denoiser(x_t)
|
||||
assert torch.allclose(result, expected)
|
||||
|
||||
|
||||
def test_denoise_step_with_prev_chunk(rtc_processor_debug_disabled):
|
||||
"""Test denoise_step with previous chunk applies guidance."""
|
||||
x_t = torch.ones(1, 20, 1)
|
||||
prev_chunk = torch.full((1, 20, 1), 0.1)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return x * 0.5
|
||||
|
||||
result = rtc_processor_debug_disabled.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
expected_result = torch.tensor(
|
||||
[
|
||||
[
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.5833],
|
||||
[1.3667],
|
||||
[1.1500],
|
||||
[0.9333],
|
||||
[0.7167],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
assert torch.allclose(result, expected_result, atol=1e-4)
|
||||
|
||||
|
||||
def test_denoise_step_adds_batch_dimension():
|
||||
"""Test denoise_step handles 2D input by adding batch dimension."""
|
||||
config = RTCConfig(execution_horizon=10, max_guidance_weight=5.0)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
# 2D input (no batch dimension)
|
||||
x_t = torch.randn(10, 6)
|
||||
prev_chunk = torch.randn(5, 6)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return x * 0.5
|
||||
|
||||
result = processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Output should be 2D (batch dimension removed)
|
||||
assert result.ndim == 2
|
||||
assert result.shape == (10, 6)
|
||||
|
||||
|
||||
def test_denoise_step_uses_custom_execution_horizon():
|
||||
"""Test denoise_step uses custom execution_horizon parameter."""
|
||||
config = RTCConfig(execution_horizon=10)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
x_t = torch.ones(1, 20, 1)
|
||||
prev_chunk = torch.full((1, 15, 1), 0.1)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return x * 0.5
|
||||
|
||||
result = processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
execution_horizon=15,
|
||||
)
|
||||
|
||||
expected_result = torch.tensor(
|
||||
[
|
||||
[
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.8000],
|
||||
[1.6818],
|
||||
[1.5636],
|
||||
[1.4455],
|
||||
[1.3273],
|
||||
[1.2091],
|
||||
[1.0909],
|
||||
[0.9727],
|
||||
[0.8545],
|
||||
[0.7364],
|
||||
[0.6182],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
assert torch.allclose(result, expected_result, atol=1e-4)
|
||||
|
||||
|
||||
def test_denoise_step_guidance_weight_at_time_zero():
|
||||
"""Test denoise_step handles time=0 (tau=1) without NaN/Inf."""
|
||||
config = RTCConfig(max_guidance_weight=10.0)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
x_t = torch.ones(1, 20, 1)
|
||||
prev_chunk = torch.full((1, 20, 1), 0.1)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return x * 0.5
|
||||
|
||||
result = processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.0),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
expected_result = torch.tensor(
|
||||
[
|
||||
[
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
[0.5000],
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
assert torch.allclose(result, expected_result, atol=1e-4)
|
||||
|
||||
|
||||
def test_denoise_step_with_real_denoise_step_partial():
|
||||
"""Test denoise_step with a real denoiser."""
|
||||
config = RTCConfig(max_guidance_weight=10.0)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
batch_size = 10
|
||||
action_dim = 6
|
||||
chunk_size = 20
|
||||
|
||||
x_t = torch.ones(batch_size, chunk_size, action_dim)
|
||||
prev_chunk = torch.full((batch_size, chunk_size, action_dim), 0.1)
|
||||
|
||||
velocity_function = torch.nn.Sequential(
|
||||
torch.nn.Linear(action_dim, 1000),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(1000, 256),
|
||||
torch.nn.ReLU(),
|
||||
torch.nn.Linear(256, action_dim),
|
||||
)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return velocity_function(x)
|
||||
|
||||
result = processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
assert result.shape == (batch_size, chunk_size, action_dim)
|
||||
|
||||
|
||||
def test_denoise_step_guidance_weight_at_time_one():
|
||||
"""Test denoise_step handles time=1 (tau=0) with max_guidance_weight clamping."""
|
||||
config = RTCConfig(max_guidance_weight=10.0)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
x_t = torch.randn(1, 50, 6)
|
||||
prev_chunk = torch.randn(1, 50, 6)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return torch.ones_like(x) * 0.5
|
||||
|
||||
# Time = 1 => tau = 0, c = (1-tau)/tau = 1/0 = inf (clamped to max_guidance_weight)
|
||||
result = processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(1.0),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Should clamp to max_guidance_weight (no Inf)
|
||||
assert not torch.any(torch.isinf(result))
|
||||
|
||||
|
||||
def test_denoise_step_tracks_debug_info(rtc_processor_debug_enabled):
|
||||
"""Test denoise_step tracks debug information when enabled."""
|
||||
x_t = torch.randn(1, 50, 6)
|
||||
prev_chunk = torch.randn(1, 50, 6)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return torch.ones_like(x) * 0.5
|
||||
|
||||
rtc_processor_debug_enabled.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Should have tracked one step
|
||||
steps = rtc_processor_debug_enabled.get_all_debug_steps()
|
||||
assert len(steps) == 1
|
||||
|
||||
# Check tracked values
|
||||
step = steps[0]
|
||||
assert step.time == 0.5
|
||||
assert step.x1_t is not None
|
||||
assert step.correction is not None
|
||||
assert step.err is not None
|
||||
assert step.weights is not None
|
||||
assert step.guidance_weight is not None
|
||||
assert step.inference_delay == 5
|
||||
|
||||
|
||||
def test_denoise_step_doesnt_track_without_debug(rtc_processor_debug_disabled):
|
||||
"""Test denoise_step doesn't track when debug disabled."""
|
||||
x_t = torch.randn(1, 50, 6)
|
||||
prev_chunk = torch.randn(1, 50, 6)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return torch.ones_like(x) * 0.5
|
||||
|
||||
rtc_processor_debug_disabled.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Should not track
|
||||
steps = rtc_processor_debug_disabled.get_all_debug_steps()
|
||||
assert len(steps) == 0
|
||||
|
||||
|
||||
# ====================== Integration Tests ======================
|
||||
|
||||
|
||||
def test_denoise_step_full_workflow():
|
||||
"""Test complete denoise_step workflow."""
|
||||
config = RTCConfig(
|
||||
enabled=True,
|
||||
prefix_attention_schedule=RTCAttentionSchedule.LINEAR,
|
||||
max_guidance_weight=5.0,
|
||||
execution_horizon=10,
|
||||
debug=True,
|
||||
)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
# Simulate two denoising steps
|
||||
x_t1 = torch.randn(1, 50, 6)
|
||||
x_t2 = torch.randn(1, 50, 6)
|
||||
|
||||
def mock_denoiser(x):
|
||||
return torch.randn_like(x) * 0.1
|
||||
|
||||
# First step - no guidance
|
||||
result1 = processor.denoise_step(
|
||||
x_t=x_t1,
|
||||
prev_chunk_left_over=None,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.8),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Second step - with guidance
|
||||
result2 = processor.denoise_step(
|
||||
x_t=x_t2,
|
||||
prev_chunk_left_over=result1,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.6),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Both should complete successfully
|
||||
assert result1.shape == (1, 50, 6)
|
||||
assert result2.shape == (1, 50, 6)
|
||||
|
||||
# Should have tracked one step (second one, first had no prev_chunk)
|
||||
steps = processor.get_all_debug_steps()
|
||||
assert len(steps) == 1
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_denoise_step_with_cuda_tensors():
|
||||
"""Test denoise_step works with CUDA tensors."""
|
||||
config = RTCConfig(execution_horizon=10, max_guidance_weight=5.0)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
x_t = torch.randn(1, 50, 6, device="cuda")
|
||||
prev_chunk = torch.randn(1, 50, 6, device="cuda")
|
||||
|
||||
def mock_denoiser(x):
|
||||
return torch.ones_like(x) * 0.5
|
||||
|
||||
result = processor.denoise_step(
|
||||
x_t=x_t,
|
||||
prev_chunk_left_over=prev_chunk,
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=mock_denoiser,
|
||||
)
|
||||
|
||||
# Result should be on CUDA
|
||||
assert result.device.type == "cuda"
|
||||
assert result.shape == x_t.shape
|
||||
|
||||
|
||||
def test_denoise_step_deterministic_with_same_inputs():
|
||||
"""Test denoise_step produces same output with same inputs."""
|
||||
config = RTCConfig(execution_horizon=10, max_guidance_weight=5.0)
|
||||
processor = RTCProcessor(config)
|
||||
|
||||
torch.manual_seed(42)
|
||||
x_t = torch.randn(1, 50, 6)
|
||||
prev_chunk = torch.randn(1, 50, 6)
|
||||
|
||||
def deterministic_denoiser(x):
|
||||
return torch.ones_like(x) * 0.5
|
||||
|
||||
result1 = processor.denoise_step(
|
||||
x_t=x_t.clone(),
|
||||
prev_chunk_left_over=prev_chunk.clone(),
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=deterministic_denoiser,
|
||||
)
|
||||
|
||||
result2 = processor.denoise_step(
|
||||
x_t=x_t.clone(),
|
||||
prev_chunk_left_over=prev_chunk.clone(),
|
||||
inference_delay=5,
|
||||
time=torch.tensor(0.5),
|
||||
original_denoise_step_partial=deterministic_denoiser,
|
||||
)
|
||||
|
||||
# Should produce identical results
|
||||
assert torch.allclose(result1, result2)
|
||||
Reference in New Issue
Block a user