lerobot/src/lerobot/rewards/sarm/processor_sarm.py

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""SARM Processor for encoding images/text and generating stage+tau targets."""

from __future__ import annotations

import random
from typing import TYPE_CHECKING, Any

import numpy as np
import torch
from PIL import Image

from lerobot.utils.import_utils import (
    _faker_available,
    _pandas_available,
    _transformers_available,
    require_package,
)

if TYPE_CHECKING or _transformers_available:
    from transformers import CLIPModel, CLIPProcessor
else:
    CLIPModel = None  # type: ignore[assignment, misc]
    CLIPProcessor = None  # type: ignore[assignment, misc]

if TYPE_CHECKING or _pandas_available:
    import pandas as pd
else:
    pd = None  # type: ignore[assignment]

if TYPE_CHECKING or _faker_available:
    from faker import Faker
else:
    Faker = None  # type: ignore[assignment, misc]

from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.processor import (
    AddBatchDimensionProcessorStep,
    DeviceProcessorStep,
    NormalizerProcessorStep,
    PolicyProcessorPipeline,
    ProcessorStep,
    RenameObservationsProcessorStep,
    from_tensor_to_numpy,
    policy_action_to_transition,
    transition_to_policy_action,
)
from lerobot.rewards.sarm.configuration_sarm import SARMConfig
from lerobot.rewards.sarm.sarm_utils import (
    apply_rewind_augmentation,
    compute_absolute_indices,
    find_stage_and_tau,
    pad_state_to_max_dim,
)
from lerobot.types import EnvTransition, PolicyAction, TransitionKey
from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


class SARMEncodingProcessorStep(ProcessorStep):
    """ProcessorStep that encodes images and text with CLIP and generates stage and progress labels for SARM."""

    def __init__(
        self,
        config: SARMConfig,
        image_key: str | None = None,
        dataset_meta=None,
        dataset_stats: dict | None = None,
    ):
        require_package("transformers", extra="sarm")
        require_package("faker", extra="sarm")
        require_package("pandas", extra="dataset")
        super().__init__()
        self.config = config
        self.image_key = image_key or config.image_key
        self.dataset_meta = dataset_meta
        self.dataset_stats = dataset_stats
        self.annotation_mode = config.annotation_mode

        # Helper to create temporal proportions dict
        def make_props_dict(names, props):
            return dict(zip(names, props, strict=True)) if names and props else None

        # Sparse annotations (always needed)
        self.sparse_temporal_proportions = make_props_dict(
            config.sparse_subtask_names, config.sparse_temporal_proportions
        )
        self.sparse_subtask_names = config.sparse_subtask_names

        # Dense annotations (only for dual mode)
        self.dense_subtask_names = config.dense_subtask_names if config.uses_dual_heads else None
        self.dense_temporal_proportions = (
            make_props_dict(config.dense_subtask_names, config.dense_temporal_proportions)
            if config.uses_dual_heads
            else None
        )

        self.device = torch.device(
            self.config.device if self.config.device else "cuda" if torch.cuda.is_available() else "cpu"
        )

        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
        self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32", use_fast=True)
        self.clip_model.to(self.device)
        self.clip_model.eval()

        self.verbs = ["move", "grasp", "rotate", "push", "pull", "slide", "lift", "place"]
        self.fake = Faker()

    def _find_episode_for_frame(self, frame_idx: int) -> int:
        """Find the episode index for a given frame index."""
        for ep_idx in range(len(self.dataset_meta.episodes)):
            ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
            ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
            if ep_start <= frame_idx < ep_end:
                return ep_idx
        return 0

    def _get_episode_indices(self, frame_indices: np.ndarray, episode_index) -> np.ndarray:
        """Get episode indices for each frame index."""
        if episode_index is None:
            return np.array([self._find_episode_for_frame(int(f)) for f in frame_indices])

        episode_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(episode_index)))

        # If single episode but multiple frames, compute episode for each frame
        if len(episode_indices) == 1 and len(frame_indices) > 1:
            return np.array([self._find_episode_for_frame(int(f)) for f in frame_indices])

        return episode_indices

    def _generate_perturbed_task(self) -> str:
        """Generate a random perturbed task string for language perturbation."""
        num_words = random.randint(1, 5)
        verb = random.choice(self.verbs)
        phrase = " ".join([verb] + self.fake.words(nb=num_words))
        return phrase

    def _get_annotation_config(self, annotation_type: str) -> tuple[list[str], dict[str, float] | None]:
        """Get global subtask names and temporal proportions for an annotation type."""
        if annotation_type == "dense":
            return self.dense_subtask_names, self.dense_temporal_proportions
        return self.sparse_subtask_names, self.sparse_temporal_proportions

    def _load_episode_annotations(
        self,
        ep_idx: int,
        episodes_df: pd.DataFrame | None,
        annotation_type: str,
        global_names: list[str],
    ) -> tuple[list | None, list | None, list | None]:
        """Load subtask annotations for an episode from DataFrame."""
        # Single-stage mode: (linear progress 0→1)
        if episodes_df is None or len(global_names) == 1:
            return None, None, None

        # Resolve column name with fallback
        def col(suffix):
            prefixed = f"{annotation_type}_{suffix}"
            return prefixed if prefixed in episodes_df.columns else suffix

        col_names = col("subtask_names")
        if col_names not in episodes_df.columns or ep_idx >= len(episodes_df):
            return None, None, None

        subtask_names = episodes_df.loc[ep_idx, col_names]
        if subtask_names is None or (isinstance(subtask_names, float) and pd.isna(subtask_names)):
            return None, None, None

        return (
            subtask_names,
            episodes_df.loc[ep_idx, col("subtask_start_frames")],
            episodes_df.loc[ep_idx, col("subtask_end_frames")],
        )

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        """
        Encode images, text, and normalize states in the transition.

        Implements SARM training data preparation:
        - Applies language perturbation (20% probability)
        - Applies rewind augmentation (80% probability)
        - Generates stage+tau targets for all frames
        - Outputs lengths tensor for valid sequence masking
        """
        new_transition = transition.copy() if hasattr(transition, "copy") else dict(transition)
        observation = new_transition.get(TransitionKey.OBSERVATION)
        comp_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})

        frame_index = comp_data.get("index")
        episode_index = comp_data.get("episode_index")

        if frame_index is None:
            raise ValueError("Frame index ('index') not found in COMPLEMENTARY_DATA")
        if episode_index is None:
            raise ValueError("Episode index ('episode_index') not found in COMPLEMENTARY_DATA")

        frame_indices = np.atleast_1d(np.asarray(from_tensor_to_numpy(frame_index)))
        episode_indices = self._get_episode_indices(frame_indices, episode_index)

        image = observation.get(self.image_key)
        if isinstance(image, torch.Tensor):
            image = image.cpu().numpy()

        # If 4D (T, C, H, W) from delta_timestamps, add batch dim
        # If 3D (C, H, W) single frame, add batch and time dims
        if image.ndim == 4:
            image = image[np.newaxis, ...]  # (T, C, H, W) -> (1, T, C, H, W)
        elif image.ndim == 3:
            image = image[np.newaxis, np.newaxis, ...]  # (C, H, W) -> (1, 1, C, H, W)

        batch_size = image.shape[0]
        total_frames = image.shape[1]  # Should be 13: 9 obs + 4 rewind placeholders
        n_obs_steps = self.config.n_obs_steps
        max_rewind_steps = self.config.max_rewind_steps
        n_obs_frames = 1 + n_obs_steps  # 9 observation frames (including current)

        # Rewind augmentation
        rewind_steps = torch.zeros(batch_size, dtype=torch.int32)
        apply_rewind = self.training and random.random() < self.config.rewind_probability

        if apply_rewind and self.dataset_meta is not None:
            for b_idx, (ep_idx, frame_idx) in enumerate(
                zip(episode_indices.tolist(), frame_indices.tolist(), strict=True)
            ):
                ep_idx, frame_idx = int(ep_idx), int(frame_idx)
                ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]

                rewind_step, _ = apply_rewind_augmentation(
                    frame_idx, ep_start, n_obs_steps, max_rewind_steps, frame_gap=self.config.frame_gap
                )
                rewind_steps[b_idx] = rewind_step

        # Compute valid lengths: n_obs_frames + rewind_steps
        lengths = n_obs_frames + rewind_steps  # (B,)

        # Apply rewind masking to images
        # For frames beyond valid length, we mask with zeros (or copy last valid frame)
        for b_idx in range(batch_size):
            valid_len = lengths[b_idx].item()
            if valid_len < total_frames:
                image[b_idx, valid_len:] = 0  # Zero out frames beyond valid length

        # Encode images with CLIP
        video_features = self._encode_images_batch(image)
        observation["video_features"] = video_features

        state_key = self.config.state_key
        state_data = observation.get(state_key)

        if isinstance(state_data, torch.Tensor):
            state_tensor = state_data.float()
        else:
            state_tensor = torch.tensor(state_data, dtype=torch.float32)

        if state_tensor.ndim == 2:
            state_tensor = state_tensor.unsqueeze(0)  # (T, D) -> (1, T, D)
        elif state_tensor.ndim == 1:
            state_tensor = state_tensor.unsqueeze(0).unsqueeze(0)  # (D,) -> (1, 1, D)

        # Apply same rewind masking to state
        for b_idx in range(batch_size):
            valid_len = lengths[b_idx].item()
            if valid_len < state_tensor.shape[1]:
                state_tensor[b_idx, valid_len:] = 0  # Zero out frames beyond valid length

        observation["state_features"] = pad_state_to_max_dim(state_tensor, self.config.max_state_dim)

        task = comp_data.get("task")
        if isinstance(task, list):
            task = task[0] if task else ""

        # Apply language perturbation during training (20% probability)
        # When perturbed, targets will be zeroed to train model to output low values for irrelevant text
        apply_perturbation = self.training and random.random() < self.config.language_perturbation_probability
        if apply_perturbation:
            task = self._generate_perturbed_task()

        # Encode text with CLIP
        observation["text_features"] = self._encode_text_clip(task, batch_size)

        # Store lengths for model
        observation["lengths"] = lengths

        # When language is perturbed, targets are zero so perturbed samples don't contribute to progress loss
        if self.dataset_meta is not None:
            episodes_df = self.dataset_meta.episodes.to_pandas()

            # Generate sparse targets
            if self.sparse_temporal_proportions is not None:
                if apply_perturbation:
                    # Zero targets when language is perturbed
                    sparse_targets = torch.zeros(batch_size, total_frames, dtype=torch.float32)
                else:
                    sparse_targets = self._compute_batch_targets(
                        frame_indices, episode_indices, lengths, rewind_steps, episodes_df, "sparse"
                    )
                observation["sparse_targets"] = sparse_targets

            # Generate dense targets (for dual mode)
            if self.config.uses_dual_heads and self.dense_temporal_proportions is not None:
                if apply_perturbation:
                    # Zero targets when language is perturbed
                    dense_targets = torch.zeros(batch_size, total_frames, dtype=torch.float32)
                else:
                    dense_targets = self._compute_batch_targets(
                        frame_indices, episode_indices, lengths, rewind_steps, episodes_df, "dense"
                    )
                observation["dense_targets"] = dense_targets

        new_transition[TransitionKey.OBSERVATION] = observation
        return new_transition

    def _compute_batch_targets(
        self,
        frame_indices: np.ndarray,
        episode_indices: np.ndarray,
        lengths: torch.Tensor,
        rewind_steps: torch.Tensor,
        episodes_df: pd.DataFrame | None,
        annotation_type: str,
    ) -> torch.Tensor:
        """Compute stage+tau targets for a batch of samples."""
        batch_size = len(frame_indices)
        n_obs_steps = self.config.n_obs_steps
        max_rewind_steps = self.config.max_rewind_steps
        total_frames = 1 + n_obs_steps + max_rewind_steps
        frame_gap = self.config.frame_gap

        global_names, temporal_props = self._get_annotation_config(annotation_type)
        targets = torch.zeros(batch_size, total_frames, dtype=torch.float32)

        for b_idx in range(batch_size):
            ep_idx = int(episode_indices[b_idx])
            frame_idx = int(frame_indices[b_idx])

            ep_start = self.dataset_meta.episodes[ep_idx]["dataset_from_index"]
            ep_end = self.dataset_meta.episodes[ep_idx]["dataset_to_index"]
            ep_length = ep_end - ep_start

            subtask_names, subtask_start_frames, subtask_end_frames = self._load_episode_annotations(
                ep_idx, episodes_df, annotation_type, global_names
            )

            # Compute observation frame indices
            obs_indices, _ = compute_absolute_indices(
                frame_idx, ep_start, ep_end, n_obs_steps, frame_gap=frame_gap
            )
            obs_indices = obs_indices.tolist()

            # Compute targets for observation frames
            for t_idx, abs_idx in enumerate(obs_indices):
                rel_frame = abs_idx - ep_start
                targets[b_idx, t_idx] = find_stage_and_tau(
                    rel_frame,
                    ep_length,
                    subtask_names,
                    subtask_start_frames,
                    subtask_end_frames,
                    global_names,
                    temporal_props,
                    return_combined=True,
                )

            # Compute targets for rewind frames (if any)
            rewind_step = rewind_steps[b_idx].item()
            if rewind_step > 0:
                _, rewind_indices = apply_rewind_augmentation(
                    frame_idx,
                    ep_start,
                    n_obs_steps,
                    max_rewind_steps,
                    frame_gap=frame_gap,
                    rewind_step=rewind_step,
                )

                for r_idx, abs_idx in enumerate(rewind_indices[:rewind_step]):
                    rel_frame = max(0, abs_idx - ep_start)
                    targets[b_idx, n_obs_steps + 1 + r_idx] = find_stage_and_tau(
                        rel_frame,
                        ep_length,
                        subtask_names,
                        subtask_start_frames,
                        subtask_end_frames,
                        global_names,
                        temporal_props,
                        return_combined=True,
                    )

        return targets

    @property
    def training(self) -> bool:
        return getattr(self, "_training_mode", True)

    def train(self, mode: bool = True):
        """Set training mode for augmentation decisions."""
        self._training_mode = mode
        return self

    def eval(self):
        """Set evaluation mode (disable augmentations)."""
        return self.train(False)

    @torch.no_grad()
    def _encode_images_batch(self, images: np.ndarray) -> torch.Tensor:
        """Encode a batch of images using CLIP.

        Args:
            images: Batched images with shape: (B, T, C, H, W)

        Returns:
            Encoded feature vectors with shape (B, T, 512)
        """

        batch_size, seq_length = images.shape[0], images.shape[1]
        images = images.reshape(batch_size * seq_length, *images.shape[2:])

        num_frames = images.shape[0]
        images_list = []
        for i in range(num_frames):
            img = images[i]
            if img.shape[0] in [1, 3]:  # Channel first (C, H, W)
                img = img.transpose(1, 2, 0)

            # Handle single channel
            if img.shape[-1] == 1:
                img = np.repeat(img, 3, axis=-1)

            if img.dtype != np.uint8:
                img = (img * 255).astype(np.uint8) if img.max() <= 1.0 else img.astype(np.uint8)

            images_list.append(Image.fromarray(img))

        all_embeddings = []
        for i in range(0, num_frames, self.config.clip_batch_size):
            batch_imgs = images_list[i : i + self.config.clip_batch_size]

            inputs = self.clip_processor(images=batch_imgs, return_tensors="pt")
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            # Get image embeddings
            # transformers 5.x returns BaseModelOutputWithPooling instead of a plain tensor
            output = self.clip_model.get_image_features(**inputs)
            if not isinstance(output, torch.Tensor):
                output = output.pooler_output
                if output is None:
                    raise ValueError("pooler_output should not be None for CLIP models.")
            embeddings = output.detach().cpu()

            # Handle single frame case
            if embeddings.dim() == 1:
                embeddings = embeddings.unsqueeze(0)

            all_embeddings.append(embeddings)

        all_embeddings = torch.cat(all_embeddings)  # (B*T, 512)
        all_embeddings = all_embeddings.reshape(batch_size, seq_length, -1)  # (B, T, 512)

        return all_embeddings

    @torch.no_grad()
    def _encode_text_clip(self, text: str, batch_size: int) -> torch.Tensor:
        """Encode text using CLIP text encoder (per SARM paper A.4).

        Args:
            text: Task description text to encode
            batch_size: Batch size to replicate for

        Returns:
            Encoded text features with shape (B, 512)
        """
        inputs = self.clip_processor.tokenizer([text], return_tensors="pt", padding=True, truncation=True)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        # transformers 5.x returns BaseModelOutputWithPooling instead of a plain tensor
        output = self.clip_model.get_text_features(**inputs)
        if not isinstance(output, torch.Tensor):
            output = output.pooler_output
            if output is None:
                raise ValueError("pooler_output should not be None for CLIP models.")
        text_embedding = output.detach().cpu()
        text_embedding = text_embedding.expand(batch_size, -1)

        return text_embedding

    def transform_features(
        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
        """Add encoded features to the observation features."""
        features[PipelineFeatureType.OBSERVATION]["video_features"] = PolicyFeature(
            type=FeatureType.VISUAL, shape=(self.config.num_frames, self.config.image_dim)
        )
        features[PipelineFeatureType.OBSERVATION]["text_features"] = PolicyFeature(
            type=FeatureType.LANGUAGE, shape=(self.config.text_dim,)
        )
        features[PipelineFeatureType.OBSERVATION]["state_features"] = PolicyFeature(
            type=FeatureType.STATE, shape=(self.config.num_frames, self.config.max_state_dim)
        )
        return features


def make_sarm_pre_post_processors(
    config: SARMConfig,
    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
    dataset_meta=None,
) -> tuple[
    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
    PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
    """Create pre-processor and post-processor pipelines for SARM."""
    return (
        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
            steps=[
                AddBatchDimensionProcessorStep(),
                RenameObservationsProcessorStep(rename_map={}),
                NormalizerProcessorStep(
                    features={**config.input_features, **config.output_features},
                    norm_map=config.normalization_mapping,
                    stats=dataset_stats,
                ),
                SARMEncodingProcessorStep(
                    config=config, dataset_meta=dataset_meta, dataset_stats=dataset_stats
                ),
                DeviceProcessorStep(device=config.device),
            ],
            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
        ),
        PolicyProcessorPipeline[PolicyAction, PolicyAction](
            steps=[DeviceProcessorStep(device="cpu")],
            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
            to_transition=policy_action_to_transition,
            to_output=transition_to_policy_action,
        ),
    )