admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-25 18:56:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

rename action_dim, state_dim to max_action_dim, max_state_dim

Browse Source
This commit is contained in:
Pepijn
2025-09-17 16:34:07 +02:00
parent 8c0cdb00a6
commit 53577f5f1a
4 changed files with 29 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/test_pi0_pi05_hub.py
+17 -13
View File
@@ -17,12 +17,12 @@ def create_dummy_stats(config):
"""Create dummy dataset statistics for testing."""
dummy_stats = {
"observation.state": {
"mean": torch.zeros(config.state_dim),
"std": torch.ones(config.state_dim),
"mean": torch.zeros(config.max_state_dim),
"std": torch.ones(config.max_state_dim),
},
"action": {
"mean": torch.zeros(config.action_dim),
"std": torch.ones(config.action_dim),
"mean": torch.zeros(config.max_action_dim),
"std": torch.ones(config.max_action_dim),
},
}
@@ -97,8 +97,8 @@ def _test_hub_loading(model_id, model_name):
print(f" - Model type: {model_name}")
print(f" - PaliGemma variant: {policy.config.paligemma_variant}")
print(f" - Action expert variant: {policy.config.action_expert_variant}")
print(f" - Action dimension: {policy.config.action_dim}")
print(f" - State dimension: {policy.config.state_dim}")
print(f" - Action dimension: {policy.config.max_action_dim}")
print(f" - State dimension: {policy.config.max_state_dim}")
print(f" - Chunk_size: {policy.config.chunk_size}")
print(f" - Tokenizer max length: {policy.config.tokenizer_max_length}")
if model_name == "PI0.5":
@@ -185,12 +185,12 @@ def _test_hub_loading(model_id, model_name):
# Create test batch
batch = {
"observation.state": torch.randn(
batch_size, policy.config.state_dim, dtype=torch.float32, device=device
batch_size, policy.config.max_state_dim, dtype=torch.float32, device=device
),
"action": torch.randn(
batch_size,
policy.config.chunk_size,
policy.config.action_dim,
policy.config.max_action_dim,
dtype=torch.float32,
device=device,
),
@@ -282,8 +282,8 @@ def test_all_base_models_hub_loading(model_id, model_type, policy_class):
print(f" - Model type: {model_type}")
print(f" - PaliGemma variant: {policy.config.paligemma_variant}")
print(f" - Action expert variant: {policy.config.action_expert_variant}")
print(f" - Action dimension: {policy.config.action_dim}")
print(f" - State dimension: {policy.config.state_dim}")
print(f" - Action dimension: {policy.config.max_action_dim}")
print(f" - State dimension: {policy.config.max_state_dim}")
print(f" - Chunk size: {policy.config.chunk_size}")
print(f" - Tokenizer max length: {policy.config.tokenizer_max_length}")
print(f" - Device: {device}")
@@ -339,10 +339,14 @@ def test_all_base_models_hub_loading(model_id, model_type, policy_class):
batch_size = 1
batch = {
"observation.state": torch.randn(
batch_size, policy.config.state_dim, dtype=torch.float32, device=device
batch_size, policy.config.max_state_dim, dtype=torch.float32, device=device
),
"action": torch.randn(
batch_size, policy.config.chunk_size, policy.config.action_dim, dtype=torch.float32, device=device
batch_size,
policy.config.chunk_size,
policy.config.max_action_dim,
dtype=torch.float32,
device=device,
),
"task": ["Pick up the object"] * batch_size,
}
@@ -369,7 +373,7 @@ def test_all_base_models_hub_loading(model_id, model_type, policy_class):
policy.eval()
with torch.no_grad():
action = policy.select_action(batch)
expected_shape = (batch_size, policy.config.action_dim)
expected_shape = (batch_size, policy.config.max_action_dim)
assert action.shape == expected_shape, (
f"{model_id}: Expected action shape {expected_shape}, got {action.shape}"
)