mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
Apply LIBERO action decode override after loading
This commit is contained in:
@@ -469,6 +469,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
# lerobot-eval or the policy server) to the freshly built steps.
|
# lerobot-eval or the policy server) to the freshly built steps.
|
||||||
_apply_groot_step_overrides(preprocessor, preprocessor_overrides)
|
_apply_groot_step_overrides(preprocessor, preprocessor_overrides)
|
||||||
_apply_groot_step_overrides(postprocessor, postprocessor_overrides)
|
_apply_groot_step_overrides(postprocessor, postprocessor_overrides)
|
||||||
|
_apply_groot_action_decode_transform(postprocessor, config.action_decode_transform)
|
||||||
return preprocessor, postprocessor
|
return preprocessor, postprocessor
|
||||||
|
|
||||||
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
preprocessor, postprocessor = _load_groot_processor_pipelines(
|
||||||
@@ -480,6 +481,7 @@ def make_groot_pre_post_processors_from_pretrained(
|
|||||||
)
|
)
|
||||||
_reconnect_groot_relative_absolute_steps(preprocessor, postprocessor)
|
_reconnect_groot_relative_absolute_steps(preprocessor, postprocessor)
|
||||||
_reconnect_groot_n1_7_pack_decode_steps(preprocessor, postprocessor)
|
_reconnect_groot_n1_7_pack_decode_steps(preprocessor, postprocessor)
|
||||||
|
_apply_groot_action_decode_transform(postprocessor, config.action_decode_transform)
|
||||||
return preprocessor, postprocessor
|
return preprocessor, postprocessor
|
||||||
|
|
||||||
|
|
||||||
@@ -552,6 +554,20 @@ def _reconnect_groot_n1_7_pack_decode_steps(
|
|||||||
step.pack_step = pack_step
|
step.pack_step = pack_step
|
||||||
|
|
||||||
|
|
||||||
|
def _apply_groot_action_decode_transform(
|
||||||
|
postprocessor: PolicyProcessorPipeline,
|
||||||
|
action_decode_transform: str | None,
|
||||||
|
) -> None:
|
||||||
|
use_libero_transform = action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO
|
||||||
|
|
||||||
|
for step in postprocessor.steps:
|
||||||
|
if isinstance(step, GrootN17ActionDecodeStep):
|
||||||
|
step.action_decode_transform = action_decode_transform
|
||||||
|
elif isinstance(step, GrootActionUnpackUnnormalizeStep):
|
||||||
|
step.libero_gripper_action = use_libero_transform
|
||||||
|
if use_libero_transform:
|
||||||
|
step.libero_gripper_binarize = True
|
||||||
|
|
||||||
|
|
||||||
def _resolve_feature_names_from_dataset_meta(dataset_meta: Any | None, feature_key: str) -> list[str] | None:
|
def _resolve_feature_names_from_dataset_meta(dataset_meta: Any | None, feature_key: str) -> list[str] | None:
|
||||||
features = getattr(dataset_meta, "features", {}) or {}
|
features = getattr(dataset_meta, "features", {}) or {}
|
||||||
@@ -1198,6 +1214,7 @@ def make_groot_pre_post_processors(
|
|||||||
stats=padded_stats,
|
stats=padded_stats,
|
||||||
normalize_min_max=True,
|
normalize_min_max=True,
|
||||||
clip_normalized_action=True,
|
clip_normalized_action=True,
|
||||||
|
libero_gripper_action=config.action_decode_transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
action_decode_step = GrootN17ActionDecodeStep(
|
action_decode_step = GrootN17ActionDecodeStep(
|
||||||
|
|||||||
@@ -1391,6 +1391,80 @@ def test_groot_n1_7_action_decode_requires_gripper_key_for_libero_transform():
|
|||||||
step({TransitionKey.ACTION: torch.zeros(1, 1, 1)})
|
step({TransitionKey.ACTION: torch.zeros(1, 1, 1)})
|
||||||
|
|
||||||
|
|
||||||
|
def test_groot_n1_7_fallback_processors_wire_libero_transform_to_postprocessor():
|
||||||
|
config = _groot_config()
|
||||||
|
dataset_stats = {
|
||||||
|
OBS_STATE: {
|
||||||
|
"min": torch.zeros(8),
|
||||||
|
"max": torch.ones(8),
|
||||||
|
},
|
||||||
|
ACTION: {
|
||||||
|
"min": torch.zeros(7),
|
||||||
|
"max": torch.ones(7),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
_, postprocessor = make_groot_pre_post_processors(config, dataset_stats=dataset_stats)
|
||||||
|
|
||||||
|
action_decode_step = next(
|
||||||
|
step for step in postprocessor.steps if isinstance(step, GrootActionUnpackUnnormalizeStep)
|
||||||
|
)
|
||||||
|
assert action_decode_step.libero_gripper_action is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_groot_n1_7_loaded_fallback_postprocessor_honors_config_action_decode_transform(tmp_path):
|
||||||
|
input_features, output_features = _groot_features(state_dim=8, action_dim=7)
|
||||||
|
dataset_stats = {
|
||||||
|
OBS_STATE: {
|
||||||
|
"min": torch.zeros(8),
|
||||||
|
"max": torch.ones(8),
|
||||||
|
},
|
||||||
|
ACTION: {
|
||||||
|
"min": torch.zeros(7),
|
||||||
|
"max": torch.ones(7),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
disabled_config = GrootConfig(
|
||||||
|
input_features=input_features,
|
||||||
|
output_features=output_features,
|
||||||
|
device="cpu",
|
||||||
|
use_bf16=False,
|
||||||
|
action_decode_transform=None,
|
||||||
|
)
|
||||||
|
preprocessor, postprocessor = make_groot_pre_post_processors(
|
||||||
|
disabled_config,
|
||||||
|
dataset_stats=dataset_stats,
|
||||||
|
)
|
||||||
|
save_dir = tmp_path / "saved_fallback_processors"
|
||||||
|
disabled_config.save_pretrained(save_dir)
|
||||||
|
preprocessor.save_pretrained(save_dir)
|
||||||
|
postprocessor.save_pretrained(save_dir)
|
||||||
|
|
||||||
|
saved_postprocessor = json.loads((save_dir / "policy_postprocessor.json").read_text())
|
||||||
|
saved_decode_config = next(
|
||||||
|
step["config"]
|
||||||
|
for step in saved_postprocessor["steps"]
|
||||||
|
if step["registry_name"] == "groot_action_unpack_unnormalize_v2"
|
||||||
|
)
|
||||||
|
assert saved_decode_config["libero_gripper_action"] is False
|
||||||
|
|
||||||
|
enabled_config = GrootConfig(
|
||||||
|
input_features=input_features,
|
||||||
|
output_features=output_features,
|
||||||
|
device="cpu",
|
||||||
|
use_bf16=False,
|
||||||
|
action_decode_transform=GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
|
||||||
|
)
|
||||||
|
_, loaded_postprocessor = make_pre_post_processors(enabled_config, pretrained_path=str(save_dir))
|
||||||
|
action_decode_step = next(
|
||||||
|
step for step in loaded_postprocessor.steps if isinstance(step, GrootActionUnpackUnnormalizeStep)
|
||||||
|
)
|
||||||
|
|
||||||
|
assert action_decode_step.libero_gripper_action is True
|
||||||
|
output = action_decode_step({TransitionKey.ACTION: torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -1.0]])})
|
||||||
|
torch.testing.assert_close(output[TransitionKey.ACTION][0, -1], torch.tensor(1.0))
|
||||||
|
|
||||||
|
|
||||||
def test_groot_n1_7_postprocessor_converts_libero_gripper_convention():
|
def test_groot_n1_7_postprocessor_converts_libero_gripper_convention():
|
||||||
step = GrootActionUnpackUnnormalizeStep(
|
step = GrootActionUnpackUnnormalizeStep(
|
||||||
env_action_dim=7,
|
env_action_dim=7,
|
||||||
|
|||||||
Reference in New Issue
Block a user