| 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 |
|
|
