src/openpi_value/models_pytorch/preprocessing_pytorch.py

from collections.abc import Sequence
import logging

import torch

from openpi_value.shared import image_tools
import openpi_value.transforms as _transforms
import kornia.augmentation as K

logger = logging.getLogger("openpi")

# Constants moved from model.py


IMAGE_RESOLUTION = (224, 224)


def preprocess_observation_pytorch(
    observation,
    *,
    train: bool = False,
    # image_keys: Sequence[str] = IMAGE_KEYS,
    image_keys: Sequence[str] = None,
    image_resolution: tuple[int, int] = IMAGE_RESOLUTION,
    return_full_obs: bool = False,
    apply_shape_visual_aug: bool = False,
    apply_blur_visual_aug: bool = False,
    p_mask_base: float = 0.0,
    state_noise_snr: float | None = None,
):
    """Torch.compile-compatible version of preprocess_observation_pytorch with simplified type annotations.

    This function avoids complex type annotations that can cause torch.compile issues.
    """

    assert image_keys is None, "Deprecated: cannot use image_key anymore"
    
    # assert not (apply_blur_visual_aug and apply_shape_visual_aug), "Cannot apply both custom and official visual augmentations"

    batch_shape = observation.state.shape[:-1]

    image_keys = list(observation.images.keys())

    part_order = {'base': 0, 'left_wrist': 1, 'right_wrist': 2}
    def simple_sort_key(k):
        part, timestep_str, _ = k.rsplit('_', 2)
        timestep = int(timestep_str)
        return (timestep, part_order[part])
    image_keys = sorted(image_keys, key=simple_sort_key)

    out_images = {}

    for key in image_keys:
        image = observation.images[key]

        # Handle both [B, C, H, W] and [B, H, W, C] formats
        is_channels_first = image.shape[1] == 3  # Check if channels are in dimension 1

        if is_channels_first:
            # Convert [B, C, H, W] to [B, H, W, C] for processing
            image = image.permute(0, 2, 3, 1)

        if image.shape[1:3] != image_resolution:
            logger.info(f"Resizing image {key} from {image.shape[1:3]} to {image_resolution}")
            image = image_tools.resize_with_pad_torch(image, *image_resolution)

        if train:
            # Convert from [-1, 1] to [0, 1] for PyTorch augmentations
            image = image / 2.0 + 0.5
            
            # Apply PyTorch-based augmentations
            if "wrist" not in key and apply_shape_visual_aug:
                # Geometric augmentations for non-wrist cameras
                height, width = image.shape[1:3]

                # Random crop and resize
                crop_height = int(height * 0.95)
                crop_width = int(width * 0.95)

                # Random crop
                max_h = height - crop_height
                max_w = width - crop_width
                if max_h > 0 and max_w > 0:
                    # Use tensor operations instead of .item() for torch.compile compatibility
                    start_h = torch.randint(0, max_h + 1, (1,), device=image.device)
                    start_w = torch.randint(0, max_w + 1, (1,), device=image.device)
                    image = image[:, start_h : start_h + crop_height, start_w : start_w + crop_width, :]

                # Resize back to original size
                image = torch.nn.functional.interpolate(
                    image.permute(0, 3, 1, 2),  # [b, h, w, c] -> [b, c, h, w]
                    size=(height, width),
                    mode="bilinear",
                    align_corners=False,
                ).permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]

                # Random rotation (small angles)
                # Use tensor operations instead of .item() for torch.compile compatibility
                angle = torch.rand(1, device=image.device) * 10 - 5  # Random angle between -5 and 5 degrees
                if torch.abs(angle) > 0.1:  # Only rotate if angle is significant
                    # Convert to radians
                    angle_rad = angle * torch.pi / 180.0

                    # Create rotation matrix
                    cos_a = torch.cos(angle_rad)
                    sin_a = torch.sin(angle_rad)

                    # Apply rotation using grid_sample
                    grid_x = torch.linspace(-1, 1, width, device=image.device)
                    grid_y = torch.linspace(-1, 1, height, device=image.device)

                    # Create meshgrid
                    grid_y, grid_x = torch.meshgrid(grid_y, grid_x, indexing="ij")

                    # Expand to batch dimension
                    grid_x = grid_x.unsqueeze(0).expand(image.shape[0], -1, -1)
                    grid_y = grid_y.unsqueeze(0).expand(image.shape[0], -1, -1)

                    # Apply rotation transformation
                    grid_x_rot = grid_x * cos_a - grid_y * sin_a
                    grid_y_rot = grid_x * sin_a + grid_y * cos_a

                    # Stack and reshape for grid_sample
                    grid = torch.stack([grid_x_rot, grid_y_rot], dim=-1)

                    image = torch.nn.functional.grid_sample(
                        image.permute(0, 3, 1, 2),  # [b, h, w, c] -> [b, c, h, w]
                        grid,
                        mode="bilinear",
                        padding_mode="zeros",
                        align_corners=False,
                    ).permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]

            # * add motionblur and gaussian blur
            if apply_blur_visual_aug:
                
                image_nchw = image.permute(0, 3, 1, 2).contiguous()

                aug = K.AugmentationSequential(
                    K.RandomMedianBlur(kernel_size=(3, 5), p=0.1),  # * prob too high
                    K.RandomMotionBlur(kernel_size=(3, 5), angle=35., direction=0.5, p=0.1),  # * smaller aug. Since the sensor is already blurry.
                    keepdim=True,
                )
                
                # Apply
                image_nchw = aug(image_nchw)
                
                # Permute back to [B, H, W, C]
                image = image_nchw.permute(0, 2, 3, 1).contiguous()
                
                
            # Clamp values to [0, 1]
            image = torch.clamp(image, 0, 1)

            # Back to [-1, 1]
            image = image * 2.0 - 1.0


        # Convert back to [B, C, H, W] format if it was originally channels-first
        if is_channels_first:
            image = image.permute(0, 3, 1, 2)  # [B, H, W, C] -> [B, C, H, W]

        out_images[key] = image

    out_masks = {}
    for key in out_images:
        if key not in observation.image_masks:
            # do not mask by default
            out_masks[key] = torch.ones(batch_shape, dtype=torch.bool, device=observation.state.device)
        else:
            out_masks[key] = observation.image_masks[key]
            
        if 'base' in key and train and p_mask_base > 0.0:
            # Randomly mask base images
            random_tensor = torch.rand(batch_shape, device=out_masks[key].device)
            base_mask = random_tensor > p_mask_base
            out_masks[key] = out_masks[key] & base_mask  # Combine with existing mask
 
    # * State augmentation
    
    
    # * Only for conveyor? using norm04
    state_std = [
        0.2079681158065796,
        0.7834290266036987,
        0.5441722273826599,
        0.14168238639831543,
        0.1750941127538681,
        0.15182428061962128,
        0.024107031524181366,
        0.19041913747787476,
        0.6899408102035522,
        0.4627247452735901,
        0.10430814325809479,
        0.1795605719089508,
        0.11770003288984299,
        0.03210258111357689,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0,
        0.0
    ],
    
    states = observation.state

    
    if state_noise_snr is not None:
        # 1. Calculate the noise standard deviation (sigma)
        # math.sqrt or **0.5 works fine for scalar operations here
        
        state_std = torch.tensor(state_std).to(states).reshape(1, -1)  # [1, state_dim]
       
        epsilon = 1e-6
        noise_scale = state_std / torch.sqrt(torch.tensor(10) ** (state_noise_snr / 10) + epsilon)
        noise_scale = noise_scale.expand(states.shape)  # Now noise_scale has shape [4, 32]
        
        
        # 2. Add Gaussian noise
        # torch.randn_like(states) creates N(0,1) noise on the correct device (GPU/CPU)
        # We then multiply by noise_scale to adjust the spread
        states += torch.randn_like(states) * noise_scale
        

        
 
    # Create a simple object with the required attributes instead of using the complex Observation class
    class SimpleProcessedObservation:
        def __init__(self, **kwargs):
            for key, value in kwargs.items():
                setattr(self, key, value)

    if return_full_obs:
        return SimpleProcessedObservation(
            images=out_images,
            image_masks=out_masks,
            
            state=states,
            tokenized_prompt=observation.tokenized_prompt,
            tokenized_prompt_mask=observation.tokenized_prompt_mask,

            token_ar_mask=observation.token_ar_mask,
            token_loss_mask=observation.token_loss_mask,

            action_advantage=observation.action_advantage,
            action_advantage_original=observation.action_advantage_original,
            
            frame_index=observation.frame_index,
            frame_index_progress=observation.frame_index_progress,
            is_failure_data=observation.is_failure_data,
            is_infer_data=observation.is_infer_data,
            episode_length=observation.episode_length,

            image_original=observation.image_original,
            episode_index=observation.episode_index,
            
            inferred_action=observation.inferred_action,
            noise=observation.noise,
            
        )
    else:
        # * Simplified for sampling value
        return SimpleProcessedObservation(
            images=out_images,
            image_masks=out_masks,
            state=states,
            tokenized_prompt=observation.tokenized_prompt,
            tokenized_prompt_mask=observation.tokenized_prompt_mask,
        )