mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 01:41:54 +00:00
removed n1.5 dependency
This commit is contained in:
nv-sachdevkartik
committed by
Andy Wrenn
Andy Wrenn
parent
58247ab9bc
commit
b6c910e936
@@ -29,8 +29,6 @@ from lerobot.configs import FeatureType, PolicyFeature
|
||||
from lerobot.policies.factory import make_policy_config, make_pre_post_processors
|
||||
from lerobot.policies.groot.configuration_groot import (
|
||||
GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
|
||||
GROOT_N1_5,
|
||||
GROOT_N1_5_BASE_MODEL,
|
||||
GROOT_N1_7,
|
||||
GROOT_N1_7_BASE_MODEL,
|
||||
GrootConfig,
|
||||
@@ -40,7 +38,6 @@ from lerobot.policies.groot.configuration_groot import (
|
||||
from lerobot.policies.groot.modeling_groot import GrootPolicy
|
||||
from lerobot.policies.groot.processor_groot import (
|
||||
GrootActionUnpackUnnormalizeStep,
|
||||
GrootEagleEncodeStep,
|
||||
GrootN17ActionDecodeStep,
|
||||
GrootN17PackInputsStep,
|
||||
GrootN17VLMEncodeStep,
|
||||
@@ -64,11 +61,9 @@ def _groot_features(
|
||||
)
|
||||
|
||||
|
||||
def _groot_config(model_version: str) -> GrootConfig:
|
||||
def _groot_config(model_version: str = GROOT_N1_7) -> GrootConfig:
|
||||
input_features, output_features = _groot_features(state_dim=8, action_dim=7)
|
||||
kwargs = {}
|
||||
if model_version == GROOT_N1_7:
|
||||
kwargs["action_decode_transform"] = GROOT_ACTION_DECODE_TRANSFORM_LIBERO
|
||||
kwargs = {"action_decode_transform": GROOT_ACTION_DECODE_TRANSFORM_LIBERO}
|
||||
return GrootConfig(
|
||||
model_version=model_version,
|
||||
input_features=input_features,
|
||||
@@ -347,19 +342,9 @@ class _DummyGrootModel(nn.Module):
|
||||
return {"action_pred": torch.zeros(batch_size, 40, 132, device=self.weight.device)}
|
||||
|
||||
|
||||
def test_groot_n1_5_defaults_are_preserved():
|
||||
def test_groot_defaults_use_n1_7():
|
||||
config = GrootConfig(device="cpu")
|
||||
|
||||
assert config.model_version == GROOT_N1_5
|
||||
assert config.base_model_path == GROOT_N1_5_BASE_MODEL
|
||||
assert config.max_state_dim == 64
|
||||
assert config.max_action_dim == 32
|
||||
assert len(config.action_delta_indices) == 16
|
||||
|
||||
|
||||
def test_groot_n1_7_explicit_selection_uses_n1_7_defaults():
|
||||
config = GrootConfig(model_version=GROOT_N1_7, device="cpu")
|
||||
|
||||
assert config.model_version == GROOT_N1_7
|
||||
assert config.base_model_path == GROOT_N1_7_BASE_MODEL
|
||||
assert config.max_state_dim == 132
|
||||
@@ -389,25 +374,17 @@ def test_groot_n1_7_rejects_legacy_libero_gripper_action_decode_transform(legacy
|
||||
)
|
||||
|
||||
|
||||
def test_groot_n1_5_rejects_action_decode_transform():
|
||||
with pytest.raises(ValueError, match="action_decode_transform"):
|
||||
GrootConfig(
|
||||
model_version=GROOT_N1_5,
|
||||
action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
|
||||
def test_groot_n1_7_path_requires_matching_model_version():
|
||||
with pytest.raises(ValueError, match="model_version"):
|
||||
GrootConfig(base_model_path=GROOT_N1_7_BASE_MODEL, device="cpu")
|
||||
@pytest.mark.parametrize("legacy_version", ["n1.5", "n1_5", "n15", "1.5"])
|
||||
def test_groot_rejects_n1_5_aliases(legacy_version):
|
||||
with pytest.raises(ValueError, match="Unsupported GR00T model_version"):
|
||||
GrootConfig(model_version=legacy_version, device="cpu")
|
||||
|
||||
|
||||
def test_groot_config_rejects_mismatched_n1_5_path_for_n1_7():
|
||||
with pytest.raises(ValueError, match="does not match base_model_path"):
|
||||
GrootConfig(
|
||||
model_version=GROOT_N1_7,
|
||||
base_model_path=GROOT_N1_5_BASE_MODEL,
|
||||
base_model_path="nvidia/GR00T-N1.5-3B",
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
@@ -528,7 +505,16 @@ def test_groot_n1_7_predict_action_chunk_truncates_to_checkpoint_valid_horizon(t
|
||||
def test_groot_from_pretrained_rejects_mismatched_caller_config(tmp_path):
|
||||
model_path = tmp_path / "GR00T-N1.7-local"
|
||||
model_path.mkdir()
|
||||
config = _groot_config(GROOT_N1_5)
|
||||
input_features, output_features = _groot_features(state_dim=8, action_dim=7)
|
||||
config = GrootConfig(
|
||||
model_version=GROOT_N1_7,
|
||||
base_model_path="nvidia/GR00T-N1.5-3B",
|
||||
input_features=input_features,
|
||||
output_features=output_features,
|
||||
device="cpu",
|
||||
use_bf16=False,
|
||||
action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="does not match base_model_path"):
|
||||
GrootPolicy.from_pretrained(model_path, config=config)
|
||||
@@ -1251,7 +1237,16 @@ def test_groot_from_pretrained_rejects_caller_config_mismatch_from_local_config(
|
||||
model_path = tmp_path / "local-checkpoint"
|
||||
model_path.mkdir()
|
||||
(model_path / "config.json").write_text('{"model_type": "Gr00tN1d7"}')
|
||||
config = _groot_config(GROOT_N1_5)
|
||||
input_features, output_features = _groot_features(state_dim=8, action_dim=7)
|
||||
config = GrootConfig(
|
||||
model_version=GROOT_N1_7,
|
||||
base_model_path="nvidia/GR00T-N1.5-3B",
|
||||
input_features=input_features,
|
||||
output_features=output_features,
|
||||
device="cpu",
|
||||
use_bf16=False,
|
||||
action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="does not match base_model_path"):
|
||||
GrootPolicy.from_pretrained(model_path, config=config)
|
||||
@@ -1266,20 +1261,9 @@ def test_groot_n1_7_processors_are_registered_lazily_without_external_gr00t():
|
||||
|
||||
assert GrootN17PackInputsStep in step_types
|
||||
assert GrootN17VLMEncodeStep in step_types
|
||||
assert GrootEagleEncodeStep not in step_types
|
||||
assert "gr00t" not in sys.modules
|
||||
|
||||
|
||||
def test_groot_n1_5_processors_still_use_eagle_path():
|
||||
config = _groot_config(GROOT_N1_5)
|
||||
|
||||
preprocessor, _ = make_groot_pre_post_processors(config)
|
||||
step_types = {type(step) for step in preprocessor.steps}
|
||||
|
||||
assert GrootEagleEncodeStep in step_types
|
||||
assert GrootN17VLMEncodeStep not in step_types
|
||||
|
||||
|
||||
def test_groot_n1_7_pack_inputs_preserves_per_sample_language():
|
||||
step = GrootN17PackInputsStep(
|
||||
action_horizon=2,
|
||||
@@ -1672,67 +1656,6 @@ def test_groot_n1_7_saved_processors_reload_through_factory_preserves_saved_stat
|
||||
assert unpack_step.env_action_dim == 7
|
||||
|
||||
|
||||
def test_groot_legacy_n1_5_processors_reload_with_compatibility_overrides(tmp_path):
|
||||
config = _groot_config(GROOT_N1_5)
|
||||
dataset_stats = {
|
||||
OBS_STATE: {
|
||||
"min": torch.full((8,), -1.0),
|
||||
"max": torch.full((8,), 1.0),
|
||||
},
|
||||
ACTION: {
|
||||
"min": torch.full((7,), -2.0),
|
||||
"max": torch.full((7,), 2.0),
|
||||
},
|
||||
}
|
||||
legacy_preprocessor_config = {
|
||||
"name": "policy_preprocessor",
|
||||
"steps": [
|
||||
{
|
||||
"registry_name": "groot_pack_inputs_v3",
|
||||
"config": {
|
||||
"state_horizon": 1,
|
||||
"action_horizon": 16,
|
||||
"max_state_dim": config.max_state_dim,
|
||||
"max_action_dim": config.max_action_dim,
|
||||
"language_key": "task",
|
||||
"formalize_language": False,
|
||||
"embodiment_tag": config.embodiment_tag,
|
||||
"embodiment_mapping": {"new_embodiment": 31},
|
||||
"normalize_min_max": False,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
legacy_postprocessor_config = {
|
||||
"name": "policy_postprocessor",
|
||||
"steps": [
|
||||
{
|
||||
"registry_name": "groot_action_unpack_unnormalize_v1",
|
||||
"config": {
|
||||
"env_action_dim": 0,
|
||||
"normalize_min_max": False,
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
(tmp_path / "policy_preprocessor.json").write_text(json.dumps(legacy_preprocessor_config))
|
||||
(tmp_path / "policy_postprocessor.json").write_text(json.dumps(legacy_postprocessor_config))
|
||||
|
||||
loaded_preprocessor, loaded_postprocessor = make_pre_post_processors(
|
||||
config,
|
||||
pretrained_path=str(tmp_path),
|
||||
dataset_stats=dataset_stats,
|
||||
)
|
||||
|
||||
pack_step = loaded_preprocessor.steps[0]
|
||||
unpack_step = loaded_postprocessor.steps[0]
|
||||
assert pack_step.normalize_min_max
|
||||
assert unpack_step.normalize_min_max
|
||||
assert unpack_step.env_action_dim == 7
|
||||
torch.testing.assert_close(pack_step.stats[OBS_STATE]["min"], dataset_stats[OBS_STATE]["min"])
|
||||
torch.testing.assert_close(pack_step.stats[ACTION]["max"], dataset_stats[ACTION]["max"])
|
||||
torch.testing.assert_close(unpack_step.stats[OBS_STATE]["min"], dataset_stats[OBS_STATE]["min"])
|
||||
torch.testing.assert_close(unpack_step.stats[ACTION]["max"], dataset_stats[ACTION]["max"])
|
||||
|
||||
|
||||
def test_groot_policy_selects_n1_7_model_class(monkeypatch):
|
||||
|
||||
@@ -1,444 +0,0 @@
|
||||
#!/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.
|
||||
|
||||
"""Test script to verify Groot policy integration with LeRobot vs the original implementation, only meant to be run locally!"""
|
||||
|
||||
import gc
|
||||
import os
|
||||
from copy import deepcopy
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.groot.configuration_groot import GrootConfig
|
||||
from lerobot.policies.groot.modeling_groot import GrootPolicy
|
||||
from lerobot.policies.groot.processor_groot import make_groot_pre_post_processors
|
||||
from lerobot.processor import PolicyProcessorPipeline
|
||||
from lerobot.types import PolicyAction
|
||||
|
||||
pytest.importorskip("gr00t")
|
||||
pytest.importorskip("transformers")
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
os.environ.get("CI") == "true" or os.environ.get("GITHUB_ACTIONS") == "true",
|
||||
reason="This test requires local Groot installation and is not meant for CI",
|
||||
)
|
||||
|
||||
|
||||
from gr00t.data.dataset import ModalityConfig # noqa: E402
|
||||
from gr00t.data.embodiment_tags import EmbodimentTag # noqa: E402
|
||||
from gr00t.data.transform.base import ComposedModalityTransform # noqa: E402
|
||||
from gr00t.model.policy import Gr00tPolicy # noqa: E402
|
||||
|
||||
# GR1 humanoid dimensions (from pretrained model metadata)
|
||||
# The actual GR1 robot has 44 dimensions for both state and action
|
||||
# GR00TTransform will pad state to 64 and truncate action to 32
|
||||
DUMMY_STATE_DIM = 44
|
||||
DUMMY_ACTION_DIM = 44
|
||||
DUMMY_ACTION_HORIZON = 16
|
||||
IMAGE_SIZE = 256
|
||||
DEVICE = "cpu"
|
||||
MODEL_PATH = "nvidia/GR00T-N1.5-3B"
|
||||
|
||||
GR1_BODY_PARTS = {
|
||||
"left_arm": 7,
|
||||
"left_hand": 6,
|
||||
"left_leg": 6,
|
||||
"neck": 3,
|
||||
"right_arm": 7,
|
||||
"right_hand": 6,
|
||||
"right_leg": 6,
|
||||
"waist": 3,
|
||||
}
|
||||
|
||||
|
||||
def cleanup_memory():
|
||||
"""Clean up GPU/MPS memory to prevent OOM errors between tests."""
|
||||
print("\nCleaning up memory...")
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
if torch.backends.mps.is_available():
|
||||
torch.mps.empty_cache()
|
||||
print("Memory cleanup complete.")
|
||||
|
||||
|
||||
def set_seed_all(seed: int):
|
||||
"""Set random seed for all RNG sources to ensure reproducibility."""
|
||||
import random
|
||||
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
# Set deterministic behavior
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.use_deterministic_algorithms(True, warn_only=True)
|
||||
|
||||
|
||||
def instantiate_lerobot_groot(
|
||||
from_pretrained: bool = False,
|
||||
model_path: str = MODEL_PATH,
|
||||
) -> tuple[
|
||||
GrootPolicy,
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
"""Instantiate LeRobot Groot policy with preprocessor and postprocessor."""
|
||||
if from_pretrained:
|
||||
policy = GrootPolicy.from_pretrained(
|
||||
pretrained_name_or_path=model_path,
|
||||
strict=False,
|
||||
)
|
||||
policy.config.embodiment_tag = "gr1"
|
||||
else:
|
||||
config = GrootConfig(
|
||||
base_model_path=model_path,
|
||||
n_action_steps=DUMMY_ACTION_HORIZON,
|
||||
chunk_size=DUMMY_ACTION_HORIZON,
|
||||
image_size=[IMAGE_SIZE, IMAGE_SIZE],
|
||||
device=DEVICE,
|
||||
embodiment_tag="gr1",
|
||||
)
|
||||
policy = GrootPolicy(config)
|
||||
|
||||
policy.to(DEVICE)
|
||||
policy.config.device = DEVICE
|
||||
|
||||
preprocessor, postprocessor = make_groot_pre_post_processors(
|
||||
config=policy.config,
|
||||
dataset_stats=None, # Pass None for dataset_stats to disable normalization (original GR00T doesn't normalize)
|
||||
)
|
||||
|
||||
return (policy, preprocessor, postprocessor)
|
||||
|
||||
|
||||
def instantiate_original_groot(
|
||||
from_pretrained: bool = False,
|
||||
model_path: str = MODEL_PATH,
|
||||
):
|
||||
"""Instantiate original Groot policy from NVIDIA's implementation."""
|
||||
from gr00t.data.transform.concat import ConcatTransform
|
||||
from gr00t.data.transform.state_action import StateActionToTensor
|
||||
from gr00t.data.transform.video import VideoToNumpy, VideoToTensor
|
||||
from gr00t.model.transforms import GR00TTransform
|
||||
|
||||
video_keys = ["video.ego_view"]
|
||||
state_keys = [
|
||||
"state"
|
||||
] # Important: Use single concatenated "state" key (not split body parts) to match preprocessing
|
||||
action_keys = [
|
||||
"action.left_arm",
|
||||
"action.right_arm",
|
||||
"action.left_hand",
|
||||
"action.right_hand",
|
||||
"action.left_leg",
|
||||
"action.right_leg",
|
||||
"action.neck",
|
||||
"action.waist",
|
||||
]
|
||||
language_keys = ["annotation.human.action.task_description"]
|
||||
|
||||
modality_config = {
|
||||
"video": ModalityConfig(
|
||||
delta_indices=[0], # Current frame only
|
||||
modality_keys=video_keys,
|
||||
),
|
||||
"state": ModalityConfig(
|
||||
delta_indices=[0],
|
||||
modality_keys=state_keys,
|
||||
),
|
||||
"action": ModalityConfig(
|
||||
delta_indices=list(range(DUMMY_ACTION_HORIZON)),
|
||||
modality_keys=action_keys,
|
||||
),
|
||||
"language": ModalityConfig(
|
||||
delta_indices=[0],
|
||||
modality_keys=language_keys,
|
||||
),
|
||||
}
|
||||
|
||||
modality_transform = ComposedModalityTransform(
|
||||
transforms=[
|
||||
VideoToTensor(apply_to=video_keys),
|
||||
VideoToNumpy(apply_to=video_keys), # Convert to numpy (GR00TTransform expects numpy arrays)
|
||||
# State is already a single concatenated key, so no StateActionToTensor needed
|
||||
# Convert action from numpy to tensor
|
||||
StateActionToTensor(apply_to=action_keys),
|
||||
# Concatenate only video and actions (state is already single key)
|
||||
ConcatTransform(
|
||||
video_concat_order=video_keys,
|
||||
state_concat_order=[], # Empty:state is already single key
|
||||
action_concat_order=action_keys,
|
||||
),
|
||||
GR00TTransform(
|
||||
max_state_dim=64,
|
||||
max_action_dim=32,
|
||||
state_horizon=1,
|
||||
action_horizon=DUMMY_ACTION_HORIZON,
|
||||
training=False,
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
policy = Gr00tPolicy(
|
||||
model_path=model_path,
|
||||
embodiment_tag=EmbodimentTag.GR1,
|
||||
modality_config=modality_config,
|
||||
modality_transform=modality_transform,
|
||||
device=DEVICE,
|
||||
)
|
||||
|
||||
return policy, modality_config, modality_transform
|
||||
|
||||
|
||||
def create_dummy_data(device=DEVICE):
|
||||
"""Create dummy data for testing both implementations."""
|
||||
batch_size = 2
|
||||
prompt = "Pick up the red cube and place it in the bin"
|
||||
state = torch.randn(batch_size, DUMMY_STATE_DIM, dtype=torch.float32, device=device)
|
||||
|
||||
batch = {
|
||||
"observation.state": state,
|
||||
"action": torch.randn(
|
||||
batch_size,
|
||||
DUMMY_ACTION_HORIZON,
|
||||
DUMMY_ACTION_DIM,
|
||||
dtype=torch.float32,
|
||||
device=device, # Action ground truth (for training)
|
||||
),
|
||||
"observation.images.ego_view": torch.rand(
|
||||
batch_size,
|
||||
3,
|
||||
IMAGE_SIZE,
|
||||
IMAGE_SIZE,
|
||||
dtype=torch.float32,
|
||||
device=device, # Images in [0, 1] range as expected by LeRobot
|
||||
),
|
||||
"task": [prompt for _ in range(batch_size)],
|
||||
}
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
def convert_lerobot_to_original_format(batch, modality_config):
|
||||
"""Convert LeRobot batch format to original Groot format.
|
||||
|
||||
The original Groot expects observations in this format:
|
||||
{
|
||||
"video.<camera_name>": np.ndarray (T, H, W, C) or (B, T, H, W, C)
|
||||
"state.<state_component>": np.ndarray (T, D) or (B, T, D)
|
||||
"action.<action_component>": np.ndarray (T, D) or (B, T, D)
|
||||
"annotation.<annotation_type>": str or list[str]
|
||||
}
|
||||
"""
|
||||
# Original Groot expects (T, H, W, C) format for images
|
||||
# LeRobot has (B, C, H, W) format, so we need to convert
|
||||
observation = {}
|
||||
|
||||
for img_key in ["ego_view"]:
|
||||
lerobot_key = f"observation.images.{img_key}"
|
||||
if lerobot_key in batch:
|
||||
img = batch[lerobot_key]
|
||||
# Convert from (B, C, H, W) to (B, T=1, H, W, C)
|
||||
img_np = img.permute(0, 2, 3, 1).unsqueeze(1).cpu().numpy()
|
||||
# Convert [0, 1] to [0, 255] uint8 as expected by original
|
||||
img_np = (img_np * 255).astype(np.uint8)
|
||||
observation[f"video.{img_key}"] = img_np
|
||||
|
||||
# Important: The Original's GR00TTransform expects "state" as (B, T, D), not split body parts
|
||||
if "observation.state" in batch:
|
||||
state = batch["observation.state"]
|
||||
state_np = state.unsqueeze(1).cpu().numpy() # (B, 1, D)
|
||||
observation["state"] = state_np
|
||||
|
||||
if "action" in batch:
|
||||
action = batch["action"]
|
||||
action_np = action.cpu().numpy()
|
||||
|
||||
start_idx = 0
|
||||
for part_name, part_dim in GR1_BODY_PARTS.items():
|
||||
end_idx = start_idx + part_dim
|
||||
observation[f"action.{part_name}"] = action_np[:, :, start_idx:end_idx]
|
||||
start_idx = end_idx
|
||||
|
||||
if "task" in batch:
|
||||
task_list = batch["task"]
|
||||
# GR00TTransform expects language with (B, T) shape for batched data
|
||||
# Create a (B, T=1) array where each element is the string directly
|
||||
bsz = len(task_list)
|
||||
task_array = np.empty((bsz, 1), dtype=object)
|
||||
for i in range(bsz):
|
||||
task_array[i, 0] = task_list[i] # Assign string directly to each (i, 0) position
|
||||
observation["annotation.human.action.task_description"] = task_array
|
||||
|
||||
return observation
|
||||
|
||||
|
||||
def test_groot_original_vs_lerobot_pretrained():
|
||||
"""Test Groot original implementation vs LeRobot implementation with pretrained weights."""
|
||||
print("Test: Groot Original vs LeRobot with Pretrained Weights (Inference)")
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
lerobot_policy, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_groot(
|
||||
from_pretrained=True
|
||||
)
|
||||
original_policy, modality_config, modality_transform = instantiate_original_groot(from_pretrained=True)
|
||||
|
||||
batch = create_dummy_data()
|
||||
batch_lerobot = deepcopy(batch)
|
||||
|
||||
print("\n[LeRobot] Running inference...")
|
||||
lerobot_policy.eval()
|
||||
batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
|
||||
|
||||
# Important: Reset seed immediately before inference to ensure identical RNG state
|
||||
torch.manual_seed(42)
|
||||
|
||||
with torch.no_grad():
|
||||
lerobot_actions = lerobot_policy.select_action(batch_lerobot_processed)
|
||||
|
||||
print("\n[Original] Running inference...")
|
||||
original_policy.model.eval()
|
||||
observation = convert_lerobot_to_original_format(batch, modality_config)
|
||||
original_obs_transformed = modality_transform(deepcopy(observation))
|
||||
|
||||
# Important: Reset seed immediately before inference to ensure identical RNG state
|
||||
torch.manual_seed(42)
|
||||
|
||||
with torch.no_grad():
|
||||
original_model_output = original_policy.model.get_action(original_obs_transformed)
|
||||
original_actions_raw = original_model_output["action_pred"] # [2, 16, 32]
|
||||
# Take first timestep
|
||||
original_actions = original_actions_raw[:, 0, :].to(lerobot_actions.device).to(lerobot_actions.dtype)
|
||||
|
||||
print("Action Comparison:")
|
||||
diff = lerobot_actions - original_actions
|
||||
abs_diff = torch.abs(diff)
|
||||
|
||||
for batch_idx in range(lerobot_actions.shape[0]):
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Batch {batch_idx}")
|
||||
print(f"{'=' * 60}")
|
||||
print(f"{'Idx':<5} {'LeRobot':<14} {'Original':<14} {'Difference':<14}")
|
||||
print("-" * 60)
|
||||
for action_idx in range(lerobot_actions.shape[1]):
|
||||
lr_val = lerobot_actions[batch_idx, action_idx].item()
|
||||
orig_val = original_actions[batch_idx, action_idx].item()
|
||||
diff_val = abs(lr_val - orig_val)
|
||||
sign = "+" if (lr_val - orig_val) > 0 else "-"
|
||||
print(f"{action_idx:<5} {lr_val:>13.6f} {orig_val:>13.6f} {sign}{diff_val:>12.6f}")
|
||||
|
||||
max_diff = abs_diff.max().item()
|
||||
tolerance = 0.001
|
||||
assert torch.allclose(lerobot_actions, original_actions, atol=tolerance), (
|
||||
f"Actions differ by more than tolerance ({tolerance}): max diff = {max_diff:.6f}"
|
||||
)
|
||||
print(f"\nSuccess: Actions match within tolerance ({tolerance})!")
|
||||
|
||||
del lerobot_policy, lerobot_preprocessor, lerobot_postprocessor
|
||||
del original_policy, modality_config, modality_transform
|
||||
del batch, batch_lerobot, observation
|
||||
cleanup_memory()
|
||||
|
||||
|
||||
def test_groot_forward_pass_comparison():
|
||||
"""Test forward pass comparison between LeRobot and Original Groot implementations."""
|
||||
print("Test: Forward Pass Comparison (Training Mode)")
|
||||
|
||||
set_seed_all(42)
|
||||
|
||||
lerobot_policy, lerobot_preprocessor, lerobot_postprocessor = instantiate_lerobot_groot(
|
||||
from_pretrained=True
|
||||
)
|
||||
original_policy, modality_config, modality_transform = instantiate_original_groot(from_pretrained=True)
|
||||
|
||||
batch = create_dummy_data()
|
||||
lerobot_policy.eval()
|
||||
original_policy.model.eval()
|
||||
|
||||
print("\n[LeRobot] Running forward pass...")
|
||||
batch_lerobot = deepcopy(batch)
|
||||
batch_lerobot_processed = lerobot_preprocessor(batch_lerobot)
|
||||
|
||||
set_seed_all(42)
|
||||
with torch.no_grad():
|
||||
lerobot_loss, lerobot_metrics = lerobot_policy.forward(batch_lerobot_processed)
|
||||
|
||||
print(f" Loss: {lerobot_loss.item():.6f}")
|
||||
|
||||
print("\n[Original] Running forward pass...")
|
||||
observation = convert_lerobot_to_original_format(batch, modality_config)
|
||||
transformed_obs = modality_transform(observation)
|
||||
|
||||
if "action" not in transformed_obs:
|
||||
action_for_forward = batch_lerobot_processed["action"]
|
||||
action_mask_for_forward = batch_lerobot_processed["action_mask"]
|
||||
|
||||
# Match action horizon if needed
|
||||
if action_for_forward.shape[1] != original_policy.model.action_horizon:
|
||||
if action_for_forward.shape[1] < original_policy.model.action_horizon:
|
||||
pad_size = original_policy.model.action_horizon - action_for_forward.shape[1]
|
||||
last_action = action_for_forward[:, -1:, :]
|
||||
padding = last_action.repeat(1, pad_size, 1)
|
||||
action_for_forward = torch.cat([action_for_forward, padding], dim=1)
|
||||
|
||||
mask_padding = torch.zeros(
|
||||
action_mask_for_forward.shape[0],
|
||||
pad_size,
|
||||
action_mask_for_forward.shape[2],
|
||||
dtype=action_mask_for_forward.dtype,
|
||||
device=action_mask_for_forward.device,
|
||||
)
|
||||
action_mask_for_forward = torch.cat([action_mask_for_forward, mask_padding], dim=1)
|
||||
else:
|
||||
action_for_forward = action_for_forward[:, : original_policy.model.action_horizon, :]
|
||||
action_mask_for_forward = action_mask_for_forward[
|
||||
:, : original_policy.model.action_horizon, :
|
||||
]
|
||||
|
||||
transformed_obs["action"] = action_for_forward
|
||||
transformed_obs["action_mask"] = action_mask_for_forward
|
||||
|
||||
set_seed_all(42)
|
||||
with torch.no_grad():
|
||||
original_outputs = original_policy.model.forward(transformed_obs)
|
||||
|
||||
original_loss = original_outputs["loss"]
|
||||
print(f" Loss: {original_loss.item():.6f}")
|
||||
|
||||
loss_diff = abs(lerobot_loss.item() - original_loss.item())
|
||||
loss_rel_diff = loss_diff / (abs(original_loss.item()) + 1e-8) * 100
|
||||
|
||||
print("\nLoss Values:")
|
||||
print(f" LeRobot: {lerobot_loss.item():.6f}")
|
||||
print(f" Original: {original_loss.item():.6f}")
|
||||
print(f" Absolute difference: {loss_diff:.6f}")
|
||||
print(f" Relative difference: {loss_rel_diff:.2f}%")
|
||||
|
||||
del lerobot_policy, lerobot_preprocessor, lerobot_postprocessor
|
||||
del original_policy, modality_config, modality_transform
|
||||
del batch, batch_lerobot, observation, transformed_obs
|
||||
cleanup_memory()
|
||||
Reference in New Issue
Block a user