labsim_transforms.py

"""
LabUtopia <-> LeRobot data transforms.

Handles converting LabUtopia observation format (camera_1_rgb, camera_2_rgb, camera_3_rgb,
state, prompt) into the LeRobot model input format, and converting model action outputs
back into the format LabUtopia expects.

This module is analogous to openpi's labsim_policy.py (LabSimInputs / LabSimOutputs).
"""

import logging
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch

logger = logging.getLogger(__name__)


def parse_image_to_uint8_hwc(image: np.ndarray) -> np.ndarray:
    """
    Parse an image to uint8 (H, W, C) format.

    LabUtopia sends images as uint8 (H, W, C), but we handle edge cases:
    - float images in [0, 1] -> convert to uint8
    - (C, H, W) layout -> transpose to (H, W, C)
    """
    image = np.asarray(image)
    if np.issubdtype(image.dtype, np.floating):
        image = (255 * np.clip(image, 0.0, 1.0)).astype(np.uint8)
    # If shape is (C, H, W) with C in {1, 3, 4}, transpose
    if image.ndim == 3 and image.shape[0] in (1, 3, 4) and image.shape[2] not in (1, 3, 4):
        image = np.transpose(image, (1, 2, 0))
    return image


def image_hwc_uint8_to_chw_float32(image: np.ndarray, normalize: bool = True) -> torch.Tensor:
    """
    Convert (H, W, C) uint8 image to (C, H, W) float32 tensor.

    Args:
        image: uint8 image in (H, W, C) format
        normalize: if True, output [0, 1] range (matches training data from video_utils);
                   if False, output [0, 255] range (for models with own image_to_float
                   step in their preprocessor, e.g. XVLA has xvla_image_to_float)
    """
    image = np.asarray(image, dtype=np.float32)
    if normalize:
        image = image / 255.0  # -> [0, 1] to match training data format
    if image.ndim == 3 and image.shape[2] in (1, 3, 4):
        image = np.transpose(image, (2, 0, 1))  # (H,W,C) -> (C,H,W)
    return torch.from_numpy(image)


def resize_image_tensor(image_tensor: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
    """
    Resize a (C, H, W) or (1, C, H, W) tensor to (target_h, target_w).
    Returns (C, H, W) tensor.
    """
    if image_tensor.ndim == 3:
        image_tensor = image_tensor.unsqueeze(0)
    resized = torch.nn.functional.interpolate(
        image_tensor.float(), size=(target_h, target_w), mode="bilinear", align_corners=False
    )
    return resized.squeeze(0)


class LabUtopiaInputTransform:
    """
    Transform LabUtopia observation dict into a LeRobot preprocessor-ready batch dict.

    LabUtopia RemoteInferenceEngine sends:
        - camera_1_rgb: (H, W, C) uint8 ndarray
        - camera_2_rgb: (H, W, C) uint8 ndarray
        - camera_3_rgb: (H, W, C) uint8 ndarray  (arm camera, optional)
        - state: (8,) float array (7 joints + 1 gripper)
        - prompt / language_instruction: str

    The LeRobot preprocessor pipeline expects a flat dict with keys like:
        - observation.images.camera_1_rgb: (C, H, W) float tensor in [0, 255]
        - observation.images.camera_2_rgb: (C, H, W) float tensor in [0, 255]
        - observation.state: (state_dim,) float tensor
        - task: str

    The preprocessor's RenameObservationsProcessorStep will rename these keys to
    match the model's feature names (e.g., camera_1_rgb -> image, camera_2_rgb -> image2).
    Then AddBatchDimension, Tokenizer, ImageToFloat, ImageNetNormalize, Device, Normalizer
    steps handle the rest.
    """

    def __init__(
        self,
        image_features: Dict[str, Tuple[int, ...]],
        camera_to_obs_key: Dict[str, str],
        state_dim: int = 8,
        normalize_images: bool = True,
    ):
        """
        Args:
            image_features: dict mapping LeRobot feature names to (C, H, W) shapes.
                e.g. {"observation.images.image": (3, 256, 256), ...}
            camera_to_obs_key: dict mapping LabUtopia camera names to the ORIGINAL
                observation key (before rename), e.g.
                {"camera_1_rgb": "observation.images.camera_1_rgb", ...}
            state_dim: dimension of the state vector.
            normalize_images: if True, normalize images to [0, 1] (for most models);
                if False, keep [0, 255] (for models like XVLA whose preprocessor
                has its own image_to_float step that divides by 255).
        """
        self.image_features = image_features
        self.camera_to_obs_key = camera_to_obs_key
        self.state_dim = state_dim
        self.normalize_images = normalize_images
        # Derive a single target (H, W) from image_features when every visual
        # feature declares the same spatial size. If shapes are missing or
        # inconsistent, leave _target_hw = None and skip resizing rather than
        # guessing one ambiguously.
        self._target_hw: Optional[Tuple[int, int]] = None
        hw_set = set()
        for shape in (image_features or {}).values():
            if isinstance(shape, (tuple, list)) and len(shape) == 3:
                hw_set.add((int(shape[1]), int(shape[2])))
        if len(hw_set) == 1:
            self._target_hw = next(iter(hw_set))

    def __call__(self, obs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Convert a raw LabUtopia observation to a LeRobot preprocessor input batch.

        Returns a dict with keys like "observation.images.camera_1_rgb",
        "observation.state", "task" — ready to be fed into the preprocessor pipeline.

        Images are resized to the target (H, W) from `image_features` (when
        uniform) and the state vector is padded/truncated to `state_dim`. Both
        are no-ops when the inputs already match the trained resolution / width.
        """
        batch = {}

        # ---- State (pad/truncate/validate to state_dim) ----
        state = np.asarray(obs.get("state", np.zeros(self.state_dim)), dtype=np.float32)
        state = np.atleast_1d(state).reshape(-1)
        if state.shape[0] != self.state_dim:
            logger.warning(
                "LabUtopiaInputTransform: state dim %d != expected state_dim %d; "
                "padding/truncating to match the trained contract.",
                state.shape[0], self.state_dim,
            )
            fitted = np.zeros(self.state_dim, dtype=np.float32)
            n = min(state.shape[0], self.state_dim)
            fitted[:n] = state[:n]
            state = fitted
        batch["observation.state"] = torch.from_numpy(state).float()

        # ---- Images (resize to image_features target H,W when known) ----
        for lab_cam_name, obs_key in self.camera_to_obs_key.items():
            if lab_cam_name in obs:
                img_hwc = parse_image_to_uint8_hwc(obs[lab_cam_name])
                # Convert to (C, H, W) float in [0, 255] range
                img_tensor = image_hwc_uint8_to_chw_float32(img_hwc, normalize=self.normalize_images)
                if self._target_hw is not None and tuple(img_tensor.shape[-2:]) != self._target_hw:
                    img_tensor = resize_image_tensor(
                        img_tensor, self._target_hw[0], self._target_hw[1]
                    )
                batch[obs_key] = img_tensor

        # ---- Language instruction / prompt ----
        prompt = obs.get("prompt", obs.get("language_instruction", ""))
        if isinstance(prompt, (np.ndarray, np.generic)):
            prompt = str(prompt.item()) if hasattr(prompt, 'item') else str(prompt)
        # Always write "task" to batch to avoid KeyError in TokenizerProcessorStep.
        # If prompt is empty, use an empty string as placeholder.
        batch["task"] = prompt if prompt else ""



        return batch