cosmos_framework/data/generator/action/transforms.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""Dataset transform wrappers for the Action project.

This module provides the ``ActionTransformPipeline`` and spatial padding utilities.

The reflection padding snaps each sample to the closest predefined resolution from
``VIDEO_RES_SIZE_INFO`` (matching VFM's approach), guaranteeing a bounded set of
output shapes that are all multiples of 16.

See :func:`~.unified_dataset.wrap_dataset` for the convenience factory that
combines datasets with transforms, and :class:`~.unified_dataset.MapToIterableAdapter`
for the map-to-iterable wrapper.
"""

from __future__ import annotations

import torch
import torchvision.transforms.functional as transforms_F

from cosmos_framework.utils import log
from cosmos_framework.data.generator.action.action_processing import (
    ActionNormalizer,
    ActionProcessor,
)
from cosmos_framework.data.generator.action.json_formatter import ActionPromptJsonFormatter
from cosmos_framework.data.generator.action.viewpoint_utils import ViewpointTextInfo
from cosmos_framework.data.generator.augmentors.duration_fps_text_timestamps import DurationFPSTextTimeStamps
from cosmos_framework.data.generator.augmentors.idle_frames_text_info import IdleFramesTextInfo
from cosmos_framework.data.generator.augmentors.resolution_text_info import ResolutionTextInfo
from cosmos_framework.data.generator.augmentors.text_tokenizer import TextTokenizerTransform
from cosmos_framework.data.generator.utils import VIDEO_RES_SIZE_INFO
from cosmos_framework.data.generator.sequence_packing import SequencePlan
from cosmos_framework.utils.generator.data_utils import get_vision_data_resolution


def _should_append_idle_frame_info(mode: object) -> bool:
    """Return whether idle-frame prompt metadata should be surfaced."""
    return mode != "inverse_dynamics"


def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]:
    """Find the closest predefined target size for a given input resolution.

    Looks up ``VIDEO_RES_SIZE_INFO[resolution]`` and selects the aspect ratio
    whose ``H/W`` ratio is closest to the input ``h/w``.

    Args:
        h: Input height in pixels.
        w: Input width in pixels.
        resolution: Resolution tier key (e.g. ``"256"``, ``"480"``, ``"720"``).

    Returns:
        ``(target_w, target_h)`` from the predefined table.

    Raises:
        ValueError: If *resolution* is not a key in ``VIDEO_RES_SIZE_INFO``.
    """
    if isinstance(resolution, int):
        resolution = str(resolution)
    if resolution not in VIDEO_RES_SIZE_INFO:
        raise ValueError(
            f"Resolution '{resolution}' not found in VIDEO_RES_SIZE_INFO. Available: {list(VIDEO_RES_SIZE_INFO.keys())}"
        )

    candidates = VIDEO_RES_SIZE_INFO[resolution]
    input_ratio = h / w

    best_key: str | None = None
    best_diff = float("inf")
    for aspect_key, (cand_w, cand_h) in candidates.items():
        cand_ratio = cand_h / cand_w
        diff = abs(input_ratio - cand_ratio)
        if diff < best_diff:
            best_diff = diff
            best_key = aspect_key

    assert best_key is not None
    target_w, target_h = candidates[best_key]
    return target_w, target_h


def reflection_pad_to_target(
    data_dict: dict,
    keys: list[str],
    keep_aspect_ratio: bool,
    target_w: int,
    target_h: int,
) -> dict:
    """Resize (aspect-preserving) and reflection-pad tensors to exact target size.

    For each key in *keys*, the tensor is:

    1. Resized so its spatial dimensions fit within ``(target_h, target_w)``
       while preserving the aspect ratio (matching VFM's
       ``ResizeLargestSideAspectPreserving``).
    2. Reflection-padded (or edge-padded when the padding exceeds the spatial
       dimension) to reach exactly ``(target_h, target_w)`` (matching VFM's
       ``ReflectionPadding``).

    After processing, the following entries are added to *data_dict*:

    - ``"image_size"``: ``torch.Tensor`` of shape ``(4,)`` containing
      ``[target_h, target_w, orig_h_resized, orig_w_resized]`` where
      ``target_h/w`` is the padded canvas size and ``orig_h/w_resized``
      is the original spatial size after aspect-preserving resize (i.e.
      the content region before padding).  After ``default_collate``
      this becomes ``(B, 4)``;  the ``IterativeJointDataLoader`` then
      splits it into per-sample ``(1, 4)`` tensors so the model can
      index as ``data_batch["image_size"][i][0][0]``.

    Args:
        data_dict: The sample dictionary (mutated in-place).
        keys: Data-dict keys whose tensors should be resized and padded.
            Tensors must have shape ``(C, H, W)`` or ``(C, T, H, W)``.
        keep_aspect_ratio: Whether to keep the aspect ratio of the input tensor.
        target_w: Target width in pixels.
        target_h: Target height in pixels.

    Returns:
        The mutated *data_dict*.
    """
    orig_h_resized: int = 0
    orig_w_resized: int = 0

    for key in keys:
        if key not in data_dict:
            continue
        tensor = data_dict[key]
        if not isinstance(tensor, torch.Tensor):
            continue

        # Extract spatial dims
        if tensor.ndim == 3:
            orig_h, orig_w = tensor.shape[-2:]
        elif tensor.ndim == 4:
            orig_h, orig_w = tensor.shape[-2:]
        else:
            raise ValueError(f"Unexpected tensor ndim={tensor.ndim} for key '{key}', expected 3 or 4")

        # Step 1: aspect-preserving resize to fit within (target_h, target_w)
        if keep_aspect_ratio:
            # Prevent upscaling the video by setting the upper bound of scaling_ratio to 1.0.
            scaling_ratio = min(target_w / orig_w, target_h / orig_h, 1.0)
            orig_h_resized = int(scaling_ratio * orig_h + 0.5)
            orig_w_resized = int(scaling_ratio * orig_w + 0.5)
            assert orig_h_resized <= target_h and orig_w_resized <= target_w, (
                f"Resize error: orig ({orig_h}, {orig_w}) target ({target_h}, {target_w}) "
                f"computed ({orig_h_resized}, {orig_w_resized})"
            )
        else:
            orig_h_resized = target_h
            orig_w_resized = target_w

        if orig_h_resized != orig_h or orig_w_resized != orig_w:
            tensor = transforms_F.resize(
                tensor,
                size=[orig_h_resized, orig_w_resized],
                interpolation=transforms_F.InterpolationMode.BICUBIC,
                antialias=True,
            )

        # Step 2: padding to exact target size (bottom and right only)
        if orig_w_resized != target_w or orig_h_resized != target_h:
            padding_right = target_w - orig_w_resized
            padding_bottom = target_h - orig_h_resized
            padding = [0, 0, padding_right, padding_bottom]

            if padding_right >= orig_w_resized or padding_bottom >= orig_h_resized:
                tensor = transforms_F.pad(tensor, padding, padding_mode="edge")
            else:
                tensor = transforms_F.pad(tensor, padding, padding_mode="reflect")

        data_dict[key] = tensor

    # image_size: shape (4,) — [target_h, target_w, orig_h_resized, orig_w_resized].
    # Matches VFM's item_dataset convention.  default_collate stacks to (B, 4);
    # IterativeJointDataLoader._get_next_sample slices to (1, 4) per sample so
    # the model can index [i][0][0].
    data_dict["image_size"] = torch.tensor(
        [target_h, target_w, orig_h_resized, orig_w_resized], dtype=torch.float
    )  # [4]

    return data_dict


def remove_reflection_padding(
    tensor: torch.Tensor,
    image_size: torch.Tensor | list[torch.Tensor] | None,
) -> torch.Tensor:
    """Remove reflection padding added by :func:`reflection_pad_to_target`.

    Content is at top-left; crops to ``(orig_h_resized, orig_w_resized)``.

    Args:
        tensor: Tensor whose last two dimensions are the padded spatial dims.
            Supports any leading dimensions, e.g. ``(C, T, H, W)`` or
            ``(C, H, W)``.
        image_size: Spatial metadata using the convention produced by
            :func:`reflection_pad_to_target`. Accepted forms are ``None`` (no
            crop), a tensor with shape ``(4,)`` or ``(1, 4)``, or a non-empty
            list whose first element has one of those tensor shapes. The four
            values are ``[target_h, target_w, orig_h_resized,
            orig_w_resized]``, where ``orig_h/w_resized`` is the original
            spatial size after aspect-preserving resize (i.e. the content
            region before padding). This matches the convention stored by
            :func:`reflection_pad_to_target` and VFM's ``ReflectionPadding``.

    Returns:
        Cropped tensor of shape ``(..., orig_h_resized, orig_w_resized)``.
    """
    if image_size is None:
        return tensor
    if isinstance(image_size, list):
        if not image_size:
            raise ValueError("Expected at least one image_size entry")
        image_size = image_size[0]  # [1,4] or [4]
    if image_size.ndim == 2 and image_size.shape[0] == 1:
        image_size = image_size[0]  # [4]
    if image_size.ndim != 1:
        raise ValueError(f"Expected image_size shape [4] or [1,4], got {tuple(image_size.shape)}")

    target_h = int(image_size[0].item())
    target_w = int(image_size[1].item())
    orig_h_resized = int(image_size[2].item())
    orig_w_resized = int(image_size[3].item())

    if orig_h_resized == target_h and orig_w_resized == target_w:
        return tensor

    return tensor[..., :orig_h_resized, :orig_w_resized].contiguous()


def build_sequence_plan_from_mode(
    mode: str,
    video_length: int,
    action_length: int,
    has_text: bool = True,
    video_temporal_downsample: int = 4,
    num_history_actions: int = 0,
) -> SequencePlan:
    """Build a SequencePlan based on the training mode.

    This function determines whether action should be included and computes the
    appropriate condition frame indexes for vision and action based on the mode.

    Args:
        mode: Training mode. One of:
            - "forward_dynamics": Predict video given first frame and all actions
            - "inverse_dynamics": Predict actions given all video frames
            - "wam": Jointly denoise video and actions given the first frame
        video_length: Number of video frames (including the conditioning frame).
        action_length: Number of action steps (typically video_length - 1).
        has_text: Whether text conditioning is available. Defaults to True.
        video_temporal_downsample: Temporal downsampling factor of the video
            tokenizer. Used to compute condition frame indexes for inverse
            dynamics mode. Defaults to 4.

    Returns:
        SequencePlan instance with appropriate settings.
        Use ``sequence_plan.has_action`` to check if action should be included.

    Raises:
        ValueError: If mode is not one of the supported modes.

    Example:
        >>> sequence_plan = build_sequence_plan_from_mode(
        ...     mode="wam",
        ...     video_length=5,
        ...     action_length=4,
        ... )
        >>> sequence_plan.has_action
        True
        >>> sequence_plan.as_dict()
        {'has_text': True, 'has_vision': True, 'has_action': True,
         'condition_frame_indexes_vision': [0], 'condition_frame_indexes_action': []}
    """
    valid_modes = ["forward_dynamics", "inverse_dynamics", "wam"]
    if mode not in valid_modes:
        raise ValueError(f"Invalid mode: {mode!r}. Must be one of {valid_modes}")

    # Determine condition frame indexes based on mode
    # forward_dynamics/wam: first frame is clean (conditioning)
    # inverse_dynamics: all frames are provided as context
    if mode in ["forward_dynamics", "wam"]:
        condition_frame_indexes_vision = [0]
    elif mode == "inverse_dynamics":
        # All frames are observed for inverse dynamics
        condition_frame_indexes_vision = list(range(0, (video_length - 1) // video_temporal_downsample + 1))
    else:
        condition_frame_indexes_vision = []

    # For action conditioning indexes:
    # forward_dynamics: all action steps are clean (conditioning)
    # inverse_dynamics/wam: action is supervised (predicted)
    # History frames (prepended) are always conditioning.
    #
    # `base_action_length` is the number of *predicted* action steps, i.e.
    # `action_length` minus the prepended history block. There are three
    # supported alignments between the action stream and the video stream
    # (`video_length` counts the leading conditioning frame at index 0):
    #
    #   Case A (dense video, action = video_length - 1)
    #       Standard forward-dynamics / WAM layout: one action per gap
    #       between consecutive video frames, no explicit initial-state
    #       action. Conditioning actions are only the prepended history.
    #
    #   Case B (dense video, action = video_length)
    #       `use_state=True` layout: one extra action is prepended to
    #       encode the initial state, so history + one initial-state
    #       action are conditioning.
    #
    #   Case C (subsampled video, general action_length > video_length)
    #       Video is subsampled by an integer factor while the action
    #       stream stays dense. `has_initial_state` mirrors Case B: it is
    #       True iff `(base_action_length - 1)` divides evenly by
    #       `(video_length - 1)`, i.e. the extra initial-state action is
    #       present.
    base_action_length = action_length - num_history_actions
    has_initial_state = video_length > 1 and (base_action_length - 1) % (video_length - 1) == 0
    if mode == "forward_dynamics":
        condition_frame_indexes_action = list(range(action_length))
    elif base_action_length == video_length - 1:
        # Case A: only the history block is conditioning.
        condition_frame_indexes_action = list(range(num_history_actions))
    elif base_action_length == video_length:
        # Case B: history block + one initial-state action.
        condition_frame_indexes_action = list(range(num_history_actions + 1))
    else:
        # Case C: same rule as A/B based on whether an initial-state action exists.
        condition_frame_indexes_action = list(
            range(num_history_actions + 1 if has_initial_state else num_history_actions)
        )

    # `action_start_frame_offset` shifts the action stream so that the first
    # non-conditioning action aligns with video frame index 1 (the first
    # frame the model must predict). The `-num_history_actions` term
    # reserves room for the history block; the additional `+1` (i.e.
    # `1 - num_history_actions`) is applied only when there is no
    # initial-state action, so the offset skips past the first action
    # slot that is instead consumed by the initial state.
    if base_action_length == video_length - 1:
        action_start_frame_offset = 1 - num_history_actions  # Case A
    elif base_action_length == video_length:
        action_start_frame_offset = -num_history_actions  # Case B
    else:
        # Case C: same rule as A/B based on whether an initial-state action exists.
        action_start_frame_offset = -num_history_actions if has_initial_state else 1 - num_history_actions

    return SequencePlan(
        has_text=has_text,
        has_vision=True,
        has_action=True,
        condition_frame_indexes_vision=condition_frame_indexes_vision,
        condition_frame_indexes_action=condition_frame_indexes_action,
        action_start_frame_offset=action_start_frame_offset,
    )


class VideoResize:
    """Resize and reflection-pad video-aligned tensors for a single sample.

    Resolution is supplied at call time. When ``resolution`` is ``None``, the
    tier is auto-detected from the sample's ``"video"`` spatial dimensions.

    Args:
        pad_keys: Data-dict keys whose values should be resized and padded.
            Pass an empty list to disable padding entirely. Defaults to
            ``["video"]``.
        keep_aspect_ratio: Whether to resize aspect-preservingly to the closest
            predefined target size before padding. Defaults to ``True``.
        log_prefix: Prefix used in debug logging.
    """

    def __init__(
        self,
        pad_keys: list[str] | None = None,
        keep_aspect_ratio: bool = True,
        log_prefix: str = "VideoResize",
    ) -> None:
        self.pad_keys = pad_keys if pad_keys is not None else ["video"]
        self.keep_aspect_ratio = keep_aspect_ratio
        self.log_prefix = log_prefix

    def __call__(self, data_dict: dict, resolution: str | int | None) -> dict:
        """Resize and pad a sample in-place.

        Args:
            data_dict: Sample dictionary containing a ``"video"`` entry.
            resolution: Resolution tier key (e.g. ``"256"``, ``"480"``,
                ``"720"``). When ``None``, auto-detected from video dimensions.

        Returns:
            The same dictionary, mutated in-place with padded tensors and an
            ``"image_size"`` entry.
        """
        video = data_dict.get("video")
        assert isinstance(video, torch.Tensor), "video is required for reflection padding"
        h, w = video.shape[-2:]

        if resolution is None:
            resolution = get_vision_data_resolution((h, w))

        if self.keep_aspect_ratio:
            target_w, target_h = find_closest_target_size(h, w, resolution)
        else:
            target_w = int(resolution)
            target_h = int(resolution)
        reflection_pad_to_target(data_dict, self.pad_keys, self.keep_aspect_ratio, target_w, target_h)

        return data_dict

    def _log_shapes(self, data_dict: dict, when: str) -> None:
        """Log tensor shapes for the configured pad keys."""
        for key in self.pad_keys:
            val = data_dict.get(key)
            if isinstance(val, torch.Tensor):
                log.debug(f"{self.log_prefix}: {when} padding '{key}' shape = {tuple(val.shape)}")


class ActionTransformPipeline:
    """A composable transform pipeline that chains ``VideoResize``, text
    tokenization, and automatic sequence plan construction.

    Reflection padding snaps each sample to the closest predefined aspect
    ratio from ``VIDEO_RES_SIZE_INFO[resolution]``, resizes
    (aspect-preserving) to fit within the target, then reflection-pads to
    the exact target size.  This guarantees a bounded set of output shapes
    (5 per resolution tier), all multiples of 16.  Resolution is supplied
    at call time via the required ``resolution`` argument to ``__call__``;
    when ``resolution`` is ``None``, the tier is auto-detected from the
    video's spatial dimensions via ``get_vision_data_resolution``.

    Text tokenization is enabled when ``tokenizer_config`` is provided.

    When the data dictionary contains a ``"mode"`` key, the pipeline automatically
    builds a ``SequencePlan`` via :func:`build_sequence_plan_from_mode` and attaches
    it as ``data_dict["sequence_plan"]``.

    Args:
        pad_keys: Data-dict keys whose values should be resized and padded. Pass
            an empty list to disable padding entirely. Defaults to ``["video"]``.
        tokenizer_config: A lazy-instantiable config dict for the VLM tokenizer. When
            ``None``, text tokenization is skipped. Defaults to ``None``.
        cfg_dropout_rate: Probability of replacing the caption with an empty string for
            classifier-free guidance. Only used when text tokenization is enabled.
            Defaults to ``0.0``.
        caption_key: The data-dict key that contains the input caption string.
            Defaults to ``"ai_caption"``.
        text_token_key: The data-dict key where tokenized text IDs will be stored.
            Defaults to ``"text_token_ids"``.
        video_temporal_downsample: Temporal downsampling factor of the video tokenizer.
            Used when building a ``SequencePlan`` for ``"inverse_dynamics"`` mode.
            Defaults to 4.
        max_action_dim: Target action dimension to pad to.  The ``"action"`` tensor
            in every sample is padded along its last dimension via
            :func:`pad_action_to_max_dim`.  Defaults to 32.
        action_channel_masking: When ``True`` (default), the original action
            dimension is stored in ``"raw_action_dim"`` so that the model masks
            loss/noise/velocity on zero-padded action channels.  When ``False``,
            ``"raw_action_dim"`` is set to ``None`` and the model treats all
            ``max_action_dim`` channels equally (original main-branch behavior).
        append_viewpoint_info: Whether to append viewpoint type metadata to the
            caption (via ``ViewpointTextInfo`` augmentor).  Requires that
            samples contain a ``"viewpoint"`` key.  Defaults to ``True``.
        append_duration_fps_timestamps: Whether to append duration and FPS metadata to the
            caption (matching VFM's ``DurationFPSTextTimeStamps`` augmentor).
            Defaults to ``True``.
        append_resolution_info: Whether to append resolution metadata to the
            caption (matching VFM's ``ResolutionTextInfo`` augmentor).
            Defaults to ``True``.
        append_idle_frames: Whether to append the idle-frame count out of the
            total action frames to the caption (Pi0.7-style metadata, via
            ``IdleFramesTextInfo`` augmentor).  The dataset is responsible for
            populating ``data_dict["idle_frames"]``; samples without it are
            silently skipped.  Idle-frame text is skipped only for
            ``"inverse_dynamics"`` mode.  Defaults to ``False`` so existing
            experiments are unaffected.
        idle_frames_dropout: Per-field dropout rate for the idle-frame segment.
            With this probability the augmentor leaves the caption unchanged
            (matching Pi0.7's ~5% per-component dropout).  Independent of the
            global ``cfg_dropout_rate``, which empties the whole caption.
            Defaults to 0.05.
        format_prompt_as_json: Whether to replace the plain text prompt with a
            structured JSON-compatible dictionary before tokenization.  When
            enabled, legacy string metadata appenders are skipped and the JSON
            formatter owns viewpoint, action, resolution, duration, FPS, and
            idle-frame fields.  Defaults to ``False``.
    """

    def __init__(
        self,
        pad_keys: list[str] | None = None,
        keep_aspect_ratio: bool = True,
        tokenizer_config: dict | None = None,
        cfg_dropout_rate: float = 0.0,
        caption_key: str = "ai_caption",
        text_token_key: str = "text_token_ids",
        video_temporal_downsample: int = 4,
        max_action_dim: int = 32,
        action_channel_masking: bool = True,
        append_viewpoint_info: bool = True,
        append_duration_fps_timestamps: bool = True,
        append_resolution_info: bool = True,
        append_idle_frames: bool = False,
        idle_frames_dropout: float = 0.05,
        format_prompt_as_json: bool = False,
    ) -> None:
        self.caption_key: str = caption_key
        self.video_temporal_downsample: int = video_temporal_downsample
        self.max_action_dim: int = max_action_dim
        self.action_channel_masking: bool = action_channel_masking
        self.action_processor: ActionProcessor = ActionProcessor(
            max_action_dim=max_action_dim,
            action_channel_masking=action_channel_masking,
        )

        # --- Spatial resize/padding stage (resolution supplied at call time) ---
        self.video_resize: VideoResize = VideoResize(
            pad_keys=pad_keys,
            keep_aspect_ratio=keep_aspect_ratio,
            log_prefix="ActionTransformPipeline",
        )
        self.pad_keys: list[str] = self.video_resize.pad_keys
        self.keep_aspect_ratio: bool = self.video_resize.keep_aspect_ratio

        self.prompt_json_formatter: ActionPromptJsonFormatter | None = None
        if format_prompt_as_json:
            self.prompt_json_formatter = ActionPromptJsonFormatter(caption_key=caption_key)

        # --- Viewpoint text augmentor (runs after ai_caption, before duration/FPS) ---
        self.viewpoint_augmentor: ViewpointTextInfo | None = None
        if append_viewpoint_info and self.prompt_json_formatter is None:
            self.viewpoint_augmentor = ViewpointTextInfo(
                input_keys=[caption_key, "viewpoint"],
                output_keys=[caption_key],
                args={"caption_key": caption_key, "viewpoint_key": "viewpoint", "enabled": True},
            )

        # --- Duration/FPS text augmentor (runs before tokenization) ---
        self.duration_fps_augmentor: DurationFPSTextTimeStamps | None = None
        if append_duration_fps_timestamps and self.prompt_json_formatter is None:
            self.duration_fps_augmentor = DurationFPSTextTimeStamps(
                input_keys=[caption_key, "video", "conditioning_fps"],
                output_keys=[caption_key],
                args={"caption_key": caption_key, "video_key": "video", "fps_key": "conditioning_fps"},
            )

        # --- Resolution text augmentor (runs before tokenization) ---
        self.resolution_info_augmentor: ResolutionTextInfo | None = None
        if append_resolution_info and self.prompt_json_formatter is None:
            self.resolution_info_augmentor = ResolutionTextInfo(
                input_keys=[caption_key, "video", "image_size"],
                output_keys=[caption_key],
                args={"caption_key": caption_key, "video_key": "video", "enabled": True},
            )

        # --- IdleFrames text augmentor (Pi0.7-style episode metadata) ---
        # Runs after resolution info, before tokenization. Per-field dropout is
        # independent from the tokenizer's global cfg_dropout_rate.
        self.idle_frames_augmentor: IdleFramesTextInfo | None = None
        if append_idle_frames and self.prompt_json_formatter is None:
            self.idle_frames_augmentor = IdleFramesTextInfo(
                input_keys=[caption_key, "idle_frames", "action"],
                output_keys=[caption_key],
                args={
                    "caption_key": caption_key,
                    "idle_frames_key": "idle_frames",
                    "action_key": "action",
                    "dropout_rate": idle_frames_dropout,
                    "enabled": True,
                },
            )

        # --- Text tokenizer augmentor ---
        self.text_tokenizer: TextTokenizerTransform | None = None
        if tokenizer_config is not None:
            self.text_tokenizer = TextTokenizerTransform(
                input_keys=[caption_key],
                output_keys=[text_token_key],
                args={
                    "tokenizer_config": tokenizer_config,
                    "cfg_dropout_rate": cfg_dropout_rate,
                },
            )

    def __call__(
        self,
        data_dict: dict,
        resolution: str | None,
        action_normalizer: ActionNormalizer | None = None,
    ) -> dict:
        """Apply the transform pipeline to a single data dictionary.

        Resolution is required at call time and is the only source of truth
        for this sample. When ``resolution`` is ``None``, the tier is
        auto-detected from the video's spatial dimensions.

        The pipeline runs in order:

        1. Resize + reflection-pad spatial dimensions to the closest
           predefined target from ``VIDEO_RES_SIZE_INFO[resolution]``.
        2. Format the caption as a structured JSON prompt (if enabled).
        3. Otherwise, append viewpoint type metadata to caption (if enabled).
        4. Append duration/FPS metadata to caption (if enabled).
        5. Append resolution metadata to caption (if enabled).
        6. Append idle-frame metadata (Pi0.7-style) to caption unless the
           sample is in inverse dynamics mode (if enabled).
        7. Tokenize caption text (if enabled).
        8. Build a ``SequencePlan`` from the ``"mode"`` key (if present).
        9. Preserve the canonical unnormalized and unpadded action as
           ``"action_raw"``, normalize real channels, pad ``"action"`` to
           ``max_action_dim``, and attach ``"action_processing_record"``.

        Args:
            data_dict: A sample dictionary as returned by a Action dataset.
            resolution: Resolution tier key (e.g. ``"256"``, ``"480"``, ``"720"``)
                for this sample. When ``None``, auto-detected from video dimensions.
            action_normalizer: Optional source-provided action normalizer. When
                present, only unpadded real action channels are normalized
                before model-space channel padding.

        Returns:
            The same dictionary, mutated in-place with padded tensors,
            ``image_size``, tokenized text IDs, a ``"sequence_plan"`` entry,
            and action processing metadata added.
        """
        mode = data_dict.get("mode")
        assert mode is not None, "mode is required"

        # 1. Resize + reflection-pad spatial dimensions to the closest predefined target from ``VIDEO_RES_SIZE_INFO[resolution]``.
        data_dict = self.video_resize(data_dict, resolution)

        # 2. Format the caption as structured JSON when requested; otherwise run the legacy string appenders.
        if self.prompt_json_formatter is not None:
            data_dict = self.prompt_json_formatter(data_dict)
        else:
            # The relabeling is only rendered by the JSON formatter.  Drop it
            # here so that enabling ``annotation_mode`` without
            # ``format_prompt_as_json`` degrades to the plain caption instead of
            # leaking a dict into the collate function.
            data_dict.pop("relabel_prompt", None)

            # 3. Append viewpoint type metadata to caption (if enabled).
            if self.viewpoint_augmentor is not None:
                result = self.viewpoint_augmentor(data_dict)
                if result is not None:
                    data_dict = result

            # 4. Append duration/FPS metadata to caption (if enabled).
            if self.duration_fps_augmentor is not None:
                result = self.duration_fps_augmentor(data_dict)
                if result is not None:
                    data_dict = result

            # 5. Append resolution metadata to caption (if enabled).
            if self.resolution_info_augmentor is not None:
                result = self.resolution_info_augmentor(data_dict)
                if result is not None:
                    data_dict = result

            # 6. Append idle-frame metadata to caption (if enabled for this mode).
            if self.idle_frames_augmentor is not None and _should_append_idle_frame_info(mode):
                result = self.idle_frames_augmentor(data_dict)
                if result is not None:
                    data_dict = result

        # 7. Tokenize caption text (if enabled).
        if self.text_tokenizer is not None:
            data_dict = self.text_tokenizer(data_dict)

        # 8. Build a ``SequencePlan`` from the ``"mode"`` key (if present).
        video = data_dict.get("video")
        action = data_dict.get("action")
        assert video is not None, "video is required"
        video_length = video.shape[1]  # [C,T,H,W] -> T
        action_length = action.shape[0] if isinstance(action, torch.Tensor) else max(video_length - 1, 0)

        # Prepend history action frames (ground-truth conditioning) if present.
        history_action = data_dict.pop("history_action", None)
        residual_gate = data_dict.get("residual_gate")
        num_history_actions = 0
        if history_action is not None and isinstance(action, torch.Tensor):
            num_history_actions = history_action.shape[0]
            action = torch.cat([history_action, action], dim=0)
            action_length += num_history_actions
            if isinstance(residual_gate, torch.Tensor):
                residual_gate = torch.cat(
                    [torch.zeros((num_history_actions, 1), dtype=residual_gate.dtype), residual_gate], dim=0
                )
                data_dict["residual_gate"] = residual_gate

        sequence_plan = build_sequence_plan_from_mode(
            mode=mode,
            video_length=video_length,
            action_length=action_length,
            video_temporal_downsample=self.video_temporal_downsample,
            num_history_actions=num_history_actions,
        )
        data_dict["sequence_plan"] = sequence_plan

        assert isinstance(action, torch.Tensor), "action tensor is required for action modes"
        data_dict = self.action_processor.preprocess_action(
            data_dict,
            action,
            action_normalizer=action_normalizer,
        )

        return data_dict