File size: 2,489 Bytes
1bfc3f0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from __future__ import annotations
from collections.abc import Mapping
import numpy as np
import torch
from ..backends.actioncodec2 import (
PresetConfig,
)
def _numpy(value):
return value.detach().cpu().float().numpy() if isinstance(value, torch.Tensor) else value
def _numpy_token_rows(token_ids):
if isinstance(token_ids, torch.Tensor):
if (
token_ids.dtype == torch.bool
or token_ids.dtype.is_floating_point
or token_ids.dtype.is_complex
):
raise ValueError("token ids must be integers")
values = token_ids.detach().cpu().long().numpy()
if values.ndim == 2:
return values.tolist()
elif isinstance(token_ids, np.ndarray) and token_ids.ndim == 2:
if token_ids.dtype.kind not in {"i", "u"}:
raise ValueError("token ids must be integers")
return np.asarray(token_ids, dtype=np.int64).tolist()
rows = []
for row in token_ids:
if isinstance(row, torch.Tensor):
if row.dtype == torch.bool or row.dtype.is_floating_point or row.dtype.is_complex:
raise ValueError("token ids must be integers")
values = row.detach().cpu().long().numpy()
else:
values = np.asarray(row)
if values.dtype.kind not in {"i", "u"} and values.size:
raise ValueError("token ids must be integers")
values = np.asarray(values, dtype=np.int64)
if values.ndim != 1:
raise ValueError(f"token id row must be 1D, got {values.shape}")
rows.append(values.tolist())
return rows
def restore_source_layout(preset: PresetConfig, decoded: Mapping[str, np.ndarray]) -> np.ndarray:
"""Place ``{component: (B,T,D_c)}`` outputs back into the preset's ``(B,T,D)`` columns."""
first = next(iter(decoded.values()))
batch, steps = first.shape[:2]
if any(values.shape[:2] != (batch, steps) for values in decoded.values()):
raise RuntimeError("decoded components disagree on batch or sequence length")
actions = np.empty((batch, steps, preset.input_dim), dtype=np.float32)
for component in preset.components:
values = decoded[component.name]
motion_width = len(component.source_dims)
actions[:, :, component.source_dims] = values[:, :, :motion_width]
if component.gripper_dim is not None:
actions[:, :, component.gripper_dim] = values[:, :, motion_width]
return actions
|