cosmos_framework/data/generator/action/action_processing.py

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

"""Shared Action processing records and normalization helpers."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Protocol

import numpy as np
import torch

from cosmos_framework.utils import log

ActionNormalizationMethod = Literal["quantile", "quantile_rot", "meanstd", "minmax"]


class ActionNormalizer(Protocol):
    """Tensor-level action normalization interface used by ActionProcessor."""

    def normalize_action(self, action: torch.Tensor) -> torch.Tensor:  # action: [...,D], returns [...,D]
        """Map raw action values into model-space values."""
        ...

    def denormalize_action(self, action: torch.Tensor) -> torch.Tensor:  # action: [...,D], returns [...,D]
        """Invert model-space action values back into raw action values."""
        ...


@dataclass(frozen=True)
class ActionAffineNormalization:
    """Resolved affine action normalizer.

    Forward normalization is ``(raw - offset) / scale``.  Inverse
    denormalization is ``normalized * scale + offset``.

    ``forward_clamp`` records lossy range-style forward clamping.  When
    ``forward_clamp_mask`` is provided, only channels with a ``True`` mask
    entry are clamped; this represents mixed UMI normalizers where some fields
    are range-clamped and others are plain affine transforms.
    """

    offset: torch.Tensor
    scale: torch.Tensor
    forward_clamp: tuple[float, float] | None = None
    forward_clamp_mask: torch.Tensor | None = None

    def normalize_action(self, action: torch.Tensor) -> torch.Tensor:  # action: [...,D], returns [...,D]
        """Normalize raw action values with resolved affine parameters."""
        offset = self.offset.to(device=action.device, dtype=action.dtype)  # [D]
        scale = self.scale.to(device=action.device, dtype=action.dtype)  # [D]
        normalized = (action - offset) / scale  # [...,D]
        if self.forward_clamp is not None:
            lo, hi = self.forward_clamp
            clamped = normalized.clamp(lo, hi)  # [...,D]
            if self.forward_clamp_mask is None:
                normalized = clamped  # [...,D]
            else:
                clamp_mask = self.forward_clamp_mask.to(device=action.device, dtype=torch.bool)  # [D]
                normalized = torch.where(clamp_mask, clamped, normalized)  # [...,D]
        return normalized  # [...,D]

    def denormalize_action(self, action: torch.Tensor) -> torch.Tensor:  # action: [...,D], returns [...,D]
        """Invert action normalization with resolved affine parameters."""
        offset = self.offset.to(device=action.device, dtype=action.dtype)  # [D]
        scale = self.scale.to(device=action.device, dtype=action.dtype)  # [D]
        return action * scale + offset  # [...,D]


def load_action_stats(stats_path: str, stats_key: str = "global") -> dict[str, np.ndarray]:
    """Load pre-computed action normalization stats from a JSON file."""
    path = Path(stats_path)
    if not path.exists():
        raise FileNotFoundError(f"Action normalization stats not found at {stats_path}.")
    log.info(f"Loading action normalization stats from {stats_path}")
    with path.open("r") as f:
        raw = json.load(f)
    stat_keys = {"mean", "std", "min", "max", "q01", "q99"}
    if stats_key in raw:
        raw = raw[stats_key]
        if not isinstance(raw, dict):
            raise TypeError(f"Action normalization stats block {stats_key!r} in {stats_path} must be a dict.")
    elif stats_key != "global" and not any(key in raw for key in stat_keys):
        raise KeyError(f"Action normalization stats block {stats_key!r} not found in {stats_path}.")
    return {k: np.array(v, dtype=np.float32) for k, v in raw.items() if k in stat_keys}


def resolve_action_normalization(
    method: ActionNormalizationMethod,
    stats: dict[str, torch.Tensor],
    apply_forward_clamp: bool = False,
) -> ActionAffineNormalization:
    """Resolve configured action stats into affine forward/inverse parameters."""
    if method == "meanstd":
        offset = stats["mean"]  # [D]
        scale = stats["std"].clamp(min=1e-8)  # [D]
        return ActionAffineNormalization(offset=offset, scale=scale)

    if method == "minmax":
        lo = stats["min"]  # [D]
        hi = stats["max"]  # [D]
    elif method in ("quantile", "quantile_rot"):
        lo = stats["q01"]  # [D]
        hi = stats["q99"]  # [D]
    else:
        raise ValueError(f"Unknown normalization method: {method!r}")

    offset = (hi + lo) / 2.0  # [D]
    scale = (hi - lo).clamp(min=1e-8) / 2.0  # [D]

    if apply_forward_clamp:
        # Ideally this hardcode should be removed, but for now we keep it so we can be aligned with mid-training checkpoints.
        forward_clamp = (-1.0, 1.0)
    else:
        forward_clamp = None

    return ActionAffineNormalization(
        offset=offset,
        scale=scale,
        forward_clamp=forward_clamp,
    )


def make_pose_action_scale_normalizer(
    action_dim: int,
    *,
    translation_scale: float = 1.0,
    rotation_scale: float = 1.0,
) -> ActionAffineNormalization:
    """Build a normalizer that maps raw pose deltas into scaled model space.

    Pose actions use the shared layout ``[translation(3), rotation(...)]``.
    The returned normalizer multiplies translation channels by
    ``translation_scale`` and rotation channels by ``rotation_scale`` during
    preprocessing, then inverts those factors during postprocessing.
    """
    if action_dim < 3:
        raise ValueError(f"Pose action_dim must be at least 3, got {action_dim}")
    if translation_scale == 0:
        raise ValueError("translation_scale must be non-zero")
    if rotation_scale == 0:
        raise ValueError("rotation_scale must be non-zero")

    offset = torch.zeros(action_dim, dtype=torch.float32)  # [D]
    scale = torch.ones(action_dim, dtype=torch.float32)  # [D]
    scale[:3] = 1.0 / float(translation_scale)  # [D]
    if action_dim > 3:
        scale[3:] = 1.0 / float(rotation_scale)  # [D]
    return ActionAffineNormalization(offset=offset, scale=scale)


@dataclass(frozen=True)
class ActionProcessingRecord:
    """Per-sample metadata needed to invert Action model-space preprocessing."""

    raw_action_dim: int
    action_normalizer: ActionNormalizer | None


def pad_action_to_max_dim(
    action: torch.Tensor, max_action_dim: int
) -> torch.Tensor:  # action: [T,D], returns [T,D_model]
    """Pad action tensor to max_action_dim along the last dimension.

    Args:
        action: Action tensor of shape (T, D) where D is the current action dimension.
        max_action_dim: Target action dimension to pad to.

    Returns:
        Padded action tensor of shape (T, max_action_dim).
    """
    if action.shape[-1] > max_action_dim:
        raise ValueError(f"Action dimension {action.shape[-1]} is greater than max_action_dim {max_action_dim}")
    if action.shape[-1] == max_action_dim:
        return action  # [T,D_model]
    padding_size = max_action_dim - action.shape[-1]
    zero_padding = torch.zeros(*action.shape[:-1], padding_size, dtype=action.dtype, device=action.device)  # [T,D_pad]
    return torch.cat([action, zero_padding], dim=-1)  # [T,D_model]


def make_batched_action_processing_fields(
    record: ActionProcessingRecord,
    batch_size: int,
    *,
    action_channel_masking: bool = True,
) -> dict[str, list[torch.Tensor | ActionProcessingRecord | None]]:
    """Build batch-list fields whose action width and inverse record cannot drift apart."""
    raw_action_dim = torch.tensor(record.raw_action_dim, dtype=torch.long) if action_channel_masking else None  # []
    return {
        "raw_action_dim": [raw_action_dim] * batch_size,
        "action_processing_record": [record] * batch_size,
    }


class ActionProcessor:
    """Build model-space actions and decode model outputs for a single sample."""

    def __init__(self, max_action_dim: int, action_channel_masking: bool = True) -> None:
        self.max_action_dim = int(max_action_dim)
        self.action_channel_masking = bool(action_channel_masking)

    def preprocess_action(
        self,
        data_dict: dict[str, Any],
        action: torch.Tensor,
        *,
        action_normalizer: ActionNormalizer | None,
    ) -> dict[str, Any]:
        """Return a sample with canonical raw and model-space action fields.

        ``action_raw`` is the exact unnormalized, unpadded action accepted by
        this method. It is the canonical external-space target for evaluation.
        ``action`` is the normalized and padded model-space representation.
        """

        raw_action_dim = int(action.shape[-1])
        if "action_raw" in data_dict:
            raise ValueError("action_raw is reserved for the canonical ActionProcessor input")
        action_raw = action.clone()  # [T,D]
        if action_normalizer is not None:
            action = action_normalizer.normalize_action(action)  # [T,D]
            if int(action.shape[-1]) != raw_action_dim:
                raise ValueError(
                    f"Action normalizer changed action width from {raw_action_dim} to {int(action.shape[-1])}"
                )

        processed_data_dict = dict(data_dict)
        processed_data_dict["action_raw"] = action_raw
        processed_data_dict["action"] = pad_action_to_max_dim(action, self.max_action_dim)  # [T,D_model]
        record = ActionProcessingRecord(
            raw_action_dim=raw_action_dim,
            action_normalizer=action_normalizer,
        )
        processed_data_dict["raw_action_dim"] = (
            torch.tensor(record.raw_action_dim, dtype=torch.long) if self.action_channel_masking else None
        )  # []
        processed_data_dict["action_processing_record"] = record
        return processed_data_dict

    @staticmethod
    def _unpad_action(action: torch.Tensor, raw_action_dim: int) -> torch.Tensor:
        """Drop model-only padded action channels."""
        if action.shape[-1] < raw_action_dim:
            raise ValueError(f"invalid raw_action_dim={raw_action_dim} for action with shape {tuple(action.shape)}")
        return action[..., :raw_action_dim].contiguous()  # [...,D_raw]

    @staticmethod
    def postprocess_action(
        action: torch.Tensor,
        record: ActionProcessingRecord,
    ) -> torch.Tensor:
        """Unpad and denormalize a generated model-space action tensor."""
        action = ActionProcessor._unpad_action(action, record.raw_action_dim)  # [...,D_raw]
        if record.action_normalizer is not None:
            action = record.action_normalizer.denormalize_action(action)  # [...,D_raw]
        return action  # [...,D_raw]


def get_action_processing_records(data_batch: dict[str, Any]) -> list[ActionProcessingRecord | None]:
    """Read all per-sample processing records from a collated Action batch."""
    records = data_batch.get("action_processing_record")
    if records is None:
        return []
    if isinstance(records, ActionProcessingRecord):
        return [records]
    if isinstance(records, list):
        for record in records:
            if record is not None and not isinstance(record, ActionProcessingRecord):
                raise TypeError(f"Unexpected action_processing_record entry type: {type(record).__name__}")
        return records
    raise TypeError(f"Unexpected action_processing_record type: {type(records).__name__}")