Guard GR00T relative action stepwise decode

This commit is contained in:
Andy Wrenn
2026-06-23 09:18:33 -07:00
parent 2ed55d2a77
commit 6126a85d60
3 changed files with 71 additions and 0 deletions
@@ -468,6 +468,14 @@ class GrootPolicy(PreTrainedPolicy):
@torch.no_grad()
def select_action(self, batch: dict[str, Tensor]) -> Tensor:
"""Select single action from action queue."""
if getattr(self.config, "use_relative_actions", False):
raise NotImplementedError(
"GrootPolicy.select_action does not support relative-action policies because cached "
"relative chunk actions can be decoded against newer observation states. Use "
"predict_action_chunk and postprocess the full chunk before queuing actions, or use "
"the RTC/chunked rollout inference path."
)
self.eval()
if len(self._action_queue) == 0:
@@ -2158,6 +2158,12 @@ class GrootN17ActionDecodeStep(ProcessorStep):
return transition
action_np = action.detach().cpu().float().numpy()
if self.use_relative_action and action_np.ndim != 3:
raise NotImplementedError(
"GrootN17ActionDecodeStep cannot decode native relative actions one step at a time. "
"Decode the full action chunk returned by predict_action_chunk while the matching "
"GrootN17PackInputsStep state is still cached, then queue the decoded absolute actions."
)
# The sync action queue postprocesses popped actions as (B, D); decode
# them as single-step (B, 1, D) chunks and squeeze the horizon back at
# the end so both ranks share the chunk decode logic below.