Reconnect GR00T relative action processors

This commit is contained in:
Andrew Wrenn
2026-06-05 09:31:04 -07:00
committed by Andy Wrenn
parent 3eb7e0b65b
commit 87e4460f60
2 changed files with 111 additions and 1 deletions
@@ -33,12 +33,14 @@ else:
ProcessorMixin = object
from lerobot.processor import (
AbsoluteActionsProcessorStep,
AddBatchDimensionProcessorStep,
DeviceProcessorStep,
PolicyAction,
PolicyProcessorPipeline,
ProcessorStep,
ProcessorStepRegistry,
RelativeActionsProcessorStep,
RenameObservationsProcessorStep,
batch_to_transition,
policy_action_to_transition,
@@ -542,9 +544,26 @@ def make_groot_pre_post_processors_from_pretrained(
to_transition=policy_action_to_transition,
to_output=transition_to_policy_action,
)
_reconnect_groot_relative_absolute_steps(preprocessor, postprocessor)
return preprocessor, postprocessor
def _reconnect_groot_relative_absolute_steps(
preprocessor: PolicyProcessorPipeline,
postprocessor: PolicyProcessorPipeline,
) -> None:
relative_step = next(
(step for step in preprocessor.steps if isinstance(step, RelativeActionsProcessorStep)),
None,
)
if relative_step is None:
return
for step in postprocessor.steps:
if isinstance(step, AbsoluteActionsProcessorStep) and step.relative_step is None:
step.relative_step = relative_step
def make_groot_pre_post_processors(
config: GrootConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
) -> tuple[