#!/usr/bin/env python

# Copyright 2024 NVIDIA Corporation and 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.

import json
from copy import copy
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any

import numpy as np
import torch
from einops import rearrange
from PIL import Image

from lerobot.utils.import_utils import _transformers_available

if TYPE_CHECKING or _transformers_available:
    from transformers import AutoProcessor, ProcessorMixin
else:
    AutoProcessor = None
    ProcessorMixin = object

from lerobot.configs import (
    FeatureType,
    NormalizationMode,
    PolicyFeature,
)
from lerobot.processor import (
    AddBatchDimensionProcessorStep,
    DeviceProcessorStep,
    PolicyAction,
    PolicyProcessorPipeline,
    ProcessorStep,
    ProcessorStepRegistry,
    RenameObservationsProcessorStep,
    batch_to_transition,
    policy_action_to_transition,
    transition_to_batch,
    transition_to_policy_action,
)
from lerobot.types import EnvTransition, TransitionKey
from lerobot.utils.constants import (
    ACTION,
    HF_LEROBOT_HOME,
    OBS_IMAGE,
    OBS_IMAGES,
    OBS_STATE,
    POLICY_POSTPROCESSOR_DEFAULT_NAME,
    POLICY_PREPROCESSOR_DEFAULT_NAME,
)

from .configuration_groot import (
    GROOT_ACTION_DECODE_TRANSFORM_LIBERO,
    GROOT_N1_7,
    GROOT_N1_7_BACKBONE_MODEL,
    GrootConfig,
    is_raw_groot_n1_7_checkpoint,
)

# Defaults for Eagle processor locations
DEFAULT_TOKENIZER_ASSETS_REPO = "lerobot/eagle2hg-processor-groot-n1p5"

N1_7_EMBODIMENT_MAPPING = {
    "oxe_droid_relative_eef_relative_joint": 24,
    "xdof_relative_eef_relative_joint": 27,
    "xdof_relative_eef_relative_joint_subtask": 27,
    "real_g1_relative_eef_relative_joints": 25,
    "real_r1_pro_sharpa_relative_eef": 26,
    "real_r1_pro_sharpa_relative_eef_human": 26,
    "real_r1_pro_sharpa_relative_eef_maxinsights": 26,
    "real_r1_pro_sharpa_relative_eef_mecka": 26,
    "unitree_g1_full_body_with_waist_height_nav_cmd": 25,
    "simpler_env_google": 0,
    "simpler_env_widowx": 1,
    "libero_sim": 2,
    "new_embodiment": 10,
}

_N1_7_RAW_STATE_CACHE: dict[str, dict[str, np.ndarray]] = {}


def _n1_7_state_cache_key(value: str | None) -> str:
    return value or "groot_n1_7_default"


@dataclass
class _GrootN17CheckpointProcessorAssets:
    """Processor metadata loaded from a raw Isaac-GR00T N1.7 checkpoint.

    Public N1.7 checkpoints store preprocessing and action-decoding choices next
    to the model weights. Keeping those values together avoids falling back to
    LeRobot defaults that are valid for older GR00T variants but change N1.7
    inputs or decoded actions.
    """

    stats: dict[str, dict[str, Any]]
    raw_stats: dict[str, Any]
    modality_config: dict[str, Any]
    embodiment_mapping: dict[str, int]
    formalize_language: bool
    valid_action_horizon: int | None
    max_action_horizon: int | None
    video_horizon: int | None
    use_percentiles: bool
    use_relative_action: bool
    clip_outliers: bool
    video_modality_keys: list[str] | None
    image_crop_size: list[int] | None
    image_target_size: list[int] | None
    shortest_image_edge: int | None
    crop_fraction: float | None
    use_albumentations: bool


def _load_n1_7_checkpoint_processor_assets(config: GrootConfig) -> _GrootN17CheckpointProcessorAssets | None:
    """Load N1.7 processor settings from checkpoint sidecar JSON files.

    Returns ``None`` for non-raw N1.7 checkpoints so the generic GR00T pipeline
    can keep using caller-provided dataset stats and config values.
    """

    if not is_raw_groot_n1_7_checkpoint(config.base_model_path):
        return None

    checkpoint_path = Path(config.base_model_path).expanduser()
    processor_config = _read_json(checkpoint_path / "processor_config.json")
    processor_kwargs = processor_config.get("processor_kwargs", {})
    if not isinstance(processor_kwargs, dict):
        processor_kwargs = {}

    all_stats = _read_json(checkpoint_path / "statistics.json")
    raw_stats = all_stats.get(config.embodiment_tag)
    if not isinstance(raw_stats, dict):
        raw_stats = {}

    modality_configs = processor_kwargs.get("modality_configs", {})
    if not isinstance(modality_configs, dict):
        modality_configs = {}
    modality_config = modality_configs.get(config.embodiment_tag)
    if not isinstance(modality_config, dict):
        modality_config = {}

    use_relative_action = bool(processor_kwargs.get("use_relative_action", False))
    stats = _load_n1_7_checkpoint_stats(
        checkpoint_path,
        processor_kwargs,
        config.embodiment_tag,
        raw_stats=raw_stats,
        modality_config=modality_config,
        use_relative_action=use_relative_action,
    )
    embodiment_mapping = _load_n1_7_embodiment_mapping(checkpoint_path) or dict(N1_7_EMBODIMENT_MAPPING)
    formalize_language = processor_kwargs.get("formalize_language", True)
    if not isinstance(formalize_language, bool):
        formalize_language = True
    clip_outliers = processor_kwargs.get("clip_outliers", True)
    if not isinstance(clip_outliers, bool):
        clip_outliers = True
    use_albumentations = processor_kwargs.get("use_albumentations", False)
    if not isinstance(use_albumentations, bool):
        use_albumentations = False

    valid_action_horizon = _load_n1_7_checkpoint_action_horizon(processor_kwargs, config.embodiment_tag)
    video_horizon = _load_n1_7_checkpoint_video_horizon(processor_kwargs, config.embodiment_tag)
    video_modality_keys = _load_n1_7_checkpoint_video_modality_keys(processor_kwargs, config.embodiment_tag)
    max_action_horizon = processor_kwargs.get("max_action_horizon")
    if not isinstance(max_action_horizon, int):
        max_action_horizon = None

    return _GrootN17CheckpointProcessorAssets(
        stats=stats,
        raw_stats=raw_stats,
        modality_config=modality_config,
        embodiment_mapping=embodiment_mapping,
        formalize_language=formalize_language,
        valid_action_horizon=valid_action_horizon,
        max_action_horizon=max_action_horizon,
        video_horizon=video_horizon,
        use_percentiles=bool(processor_kwargs.get("use_percentiles", False)),
        use_relative_action=use_relative_action,
        clip_outliers=clip_outliers,
        video_modality_keys=video_modality_keys,
        image_crop_size=_as_int_pair(processor_kwargs.get("image_crop_size")),
        image_target_size=_as_int_pair(processor_kwargs.get("image_target_size")),
        shortest_image_edge=_as_optional_int(processor_kwargs.get("shortest_image_edge")),
        crop_fraction=_as_optional_float(processor_kwargs.get("crop_fraction")),
        use_albumentations=use_albumentations,
    )


def _read_json(path: Path) -> dict[str, Any]:
    try:
        with path.open() as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}
    return data if isinstance(data, dict) else {}


def _load_n1_7_embodiment_mapping(checkpoint_path: Path) -> dict[str, int] | None:
    mapping = _read_json(checkpoint_path / "embodiment_id.json")
    if not mapping:
        return None
    parsed: dict[str, int] = {}
    for key, value in mapping.items():
        if not isinstance(key, str):
            continue
        try:
            parsed[key] = int(value)
        except (TypeError, ValueError):
            continue
    return parsed or None


def _load_n1_7_checkpoint_stats(
    checkpoint_path: Path,
    processor_kwargs: dict[str, Any],
    embodiment_tag: str,
    *,
    raw_stats: dict[str, Any] | None = None,
    modality_config: dict[str, Any] | None = None,
    use_relative_action: bool = False,
) -> dict[str, dict[str, Any]]:
    """Convert checkpoint modality-group stats into LeRobot flat tensor stats.

    Isaac-GR00T keeps statistics keyed by semantic groups such as EEF pose and
    joints. LeRobot normalizers operate over a single vector, so this function
    preserves checkpoint group order while flattening each selected statistic.
    """

    if raw_stats is None:
        all_stats = _read_json(checkpoint_path / "statistics.json")
        raw_stats = all_stats.get(embodiment_tag)
    if not isinstance(raw_stats, dict):
        return {}

    if modality_config is None:
        modality_configs = processor_kwargs.get("modality_configs", {})
        if not isinstance(modality_configs, dict):
            return {}
        modality_config = modality_configs.get(embodiment_tag)
    if not isinstance(modality_config, dict):
        return {}

    use_percentiles = processor_kwargs.get("use_percentiles", False)
    return {
        OBS_STATE: _flatten_n1_7_modality_stats(
            embodiment_stats=raw_stats,
            embodiment_config=modality_config,
            modality="state",
            use_percentiles=bool(use_percentiles),
            use_relative_action=use_relative_action,
        ),
        ACTION: _flatten_n1_7_modality_stats(
            embodiment_stats=raw_stats,
            embodiment_config=modality_config,
            modality="action",
            use_percentiles=bool(use_percentiles),
            use_relative_action=use_relative_action,
        ),
    }


def _load_n1_7_checkpoint_action_horizon(
    processor_kwargs: dict[str, Any],
    embodiment_tag: str,
) -> int | None:
    modality_configs = processor_kwargs.get("modality_configs", {})
    if not isinstance(modality_configs, dict):
        return None
    embodiment_config = modality_configs.get(embodiment_tag, {})
    if not isinstance(embodiment_config, dict):
        return None
    action_config = embodiment_config.get("action", {})
    if not isinstance(action_config, dict):
        return None
    delta_indices = action_config.get("delta_indices", [])
    if not isinstance(delta_indices, list):
        return None
    return len(delta_indices) or None


def _load_n1_7_checkpoint_video_horizon(
    processor_kwargs: dict[str, Any],
    embodiment_tag: str,
) -> int | None:
    modality_configs = processor_kwargs.get("modality_configs", {})
    if not isinstance(modality_configs, dict):
        return None
    embodiment_config = modality_configs.get(embodiment_tag, {})
    if not isinstance(embodiment_config, dict):
        return None
    video_config = embodiment_config.get("video", {})
    if not isinstance(video_config, dict):
        return None
    delta_indices = video_config.get("delta_indices", [])
    if not isinstance(delta_indices, list):
        return None
    return len(delta_indices) or None


def _load_n1_7_checkpoint_video_modality_keys(
    processor_kwargs: dict[str, Any],
    embodiment_tag: str,
) -> list[str] | None:
    modality_configs = processor_kwargs.get("modality_configs", {})
    if not isinstance(modality_configs, dict):
        return None
    embodiment_config = modality_configs.get(embodiment_tag, {})
    if not isinstance(embodiment_config, dict):
        return None
    video_config = embodiment_config.get("video", {})
    if not isinstance(video_config, dict):
        return None
    modality_keys = video_config.get("modality_keys", [])
    if not isinstance(modality_keys, list):
        return None
    keys = [key for key in modality_keys if isinstance(key, str)]
    return keys or None


def _as_int_pair(value: Any) -> list[int] | None:
    if not isinstance(value, (list, tuple)) or len(value) != 2:
        return None
    try:
        return [int(value[0]), int(value[1])]
    except (TypeError, ValueError):
        return None


def _as_optional_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _as_optional_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _flatten_n1_7_modality_stats(
    *,
    embodiment_stats: dict[str, Any],
    embodiment_config: dict[str, Any],
    modality: str,
    use_percentiles: bool,
    use_relative_action: bool,
) -> dict[str, list[float]]:
    """Flatten one N1.7 modality's grouped statistics in checkpoint order.

    When checkpoints request percentile normalization, q01/q99 replace min/max
    for regular groups. Relative action groups read from ``relative_action``
    stats and keep min/max, matching Isaac-GR00T's processor override.
    """

    source_stats = embodiment_stats.get(modality, {})
    modality_config = embodiment_config.get(modality, {})
    if not isinstance(source_stats, dict) or not isinstance(modality_config, dict):
        return {}
    modality_keys = modality_config.get("modality_keys", [])
    if not isinstance(modality_keys, list):
        return {}

    flattened: dict[str, list[float]] = {}
    action_configs = modality_config.get("action_configs", []) if modality == "action" else []
    if not isinstance(action_configs, list):
        action_configs = []
    relative_stats = embodiment_stats.get("relative_action", {})
    if not isinstance(relative_stats, dict):
        relative_stats = {}

    for stat_name in ("min", "max", "mean", "std"):
        values: list[float] = []
        source_stat_name = stat_name
        if use_percentiles and stat_name == "min":
            source_stat_name = "q01"
        elif use_percentiles and stat_name == "max":
            source_stat_name = "q99"

        for idx, modality_key in enumerate(modality_keys):
            if not isinstance(modality_key, str):
                continue
            key_source_stats = source_stats
            key_stat_name = source_stat_name
            if modality == "action" and use_relative_action and idx < len(action_configs):
                action_config = action_configs[idx]
                if isinstance(action_config, dict) and _config_value(action_config.get("rep")) == "relative":
                    key_source_stats = relative_stats
                    key_stat_name = stat_name
            key_stats = key_source_stats.get(modality_key, {})
            if not isinstance(key_stats, dict):
                raise KeyError(f"Missing statistics for {modality}.{modality_key}")
            raw_values = key_stats.get(key_stat_name)
            if raw_values is None:
                raise KeyError(f"Missing '{key_stat_name}' statistics for {modality}.{modality_key}")
            values.extend(_as_float_list(raw_values))
        if values:
            flattened[stat_name] = values

    return flattened


def _as_float_list(values: Any) -> list[float]:
    if values is None:
        return []
    if isinstance(values, torch.Tensor):
        return values.detach().cpu().reshape(-1).float().tolist()
    if isinstance(values, np.ndarray):
        return values.reshape(-1).astype(np.float32).tolist()
    if isinstance(values, (list, tuple)):
        flattened: list[float] = []
        for value in values:
            flattened.extend(_as_float_list(value))
        return flattened
    return [float(values)]


def _config_value(value: Any) -> str:
    if hasattr(value, "value"):
        value = value.value
    text = str(value).lower()
    return {
        "relative": "relative",
        "absolute": "absolute",
        "delta": "delta",
        "eef": "eef",
        "non_eef": "non_eef",
        "default": "default",
        "xyz_rot6d": "xyz+rot6d",
        "xyz+rot6d": "xyz+rot6d",
        "xyz_rotvec": "xyz+rotvec",
        "xyz+rotvec": "xyz+rotvec",
    }.get(text, text)


def _has_modality_stats(stats: dict[str, dict[str, Any]] | None) -> bool:
    if not stats:
        return False
    return any(bool(modality_stats) for modality_stats in stats.values())


def _legacy_groot_processor_overrides(
    config: GrootConfig,
    dataset_stats: dict[str, dict[str, torch.Tensor]] | None,
    preprocessor_overrides: dict[str, Any] | None = None,
    postprocessor_overrides: dict[str, Any] | None = None,
) -> tuple[dict[str, Any], dict[str, Any]]:
    """Patch older serialized Groot processors with fields current processors expect."""

    preprocessor_overrides = dict(preprocessor_overrides or {})
    postprocessor_overrides = dict(postprocessor_overrides or {})
    pack_inputs_key = (
        "groot_n1_7_pack_inputs_v1" if config.model_version == GROOT_N1_7 else "groot_pack_inputs_v3"
    )

    pack_input_overrides = dict(preprocessor_overrides.get(pack_inputs_key, {}))
    pack_input_overrides["normalize_min_max"] = True
    if dataset_stats is not None and config.model_version != GROOT_N1_7:
        pack_input_overrides["stats"] = dataset_stats
    preprocessor_overrides[pack_inputs_key] = pack_input_overrides

    try:
        env_action_dim = int(config.output_features[ACTION].shape[0])
    except Exception:
        env_action_dim = 0
    action_unpack_overrides = dict(postprocessor_overrides.get("groot_action_unpack_unnormalize_v1", {}))
    action_unpack_overrides["normalize_min_max"] = True
    action_unpack_overrides["env_action_dim"] = env_action_dim
    if dataset_stats is not None and config.model_version != GROOT_N1_7:
        action_unpack_overrides["stats"] = dataset_stats
    postprocessor_overrides["groot_action_unpack_unnormalize_v1"] = action_unpack_overrides

    return preprocessor_overrides, postprocessor_overrides


def _local_processor_config_has_step(pretrained_path: str, config_filename: str, step_name: str) -> bool:
    path = Path(pretrained_path).expanduser()
    if not path.is_dir():
        return False
    config = _read_json(path / config_filename)
    steps = config.get("steps", [])
    if not isinstance(steps, list):
        return False
    return any(isinstance(step, dict) and step.get("registry_name") == step_name for step in steps)


def make_groot_pre_post_processors_from_pretrained(
    config: GrootConfig,
    pretrained_path: str,
    *,
    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
    preprocessor_overrides: dict[str, Any] | None = None,
    postprocessor_overrides: dict[str, Any] | None = None,
    preprocessor_config_filename: str = f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
    postprocessor_config_filename: str = f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
) -> tuple[
    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
    PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
    """Load Groot processors while preserving compatibility with older serialized configs."""

    if is_raw_groot_n1_7_checkpoint(pretrained_path):
        processor_cfg = copy(config)
        processor_cfg.base_model_path = str(pretrained_path)
        return make_groot_pre_post_processors(
            config=processor_cfg,
            dataset_stats=dataset_stats,
        )

    if (
        config.model_version == GROOT_N1_7
        and _local_processor_config_has_step(
            pretrained_path,
            postprocessor_config_filename,
            "groot_n1_7_action_decode_v1",
        )
    ):
        # Converted raw N1.7 checkpoints already carry the checkpoint-specific
        # action decoder. Adding the legacy action-unpack override would target
        # a step that is not present and break loading.
        preprocessor_overrides = dict(preprocessor_overrides or {})
        postprocessor_overrides = dict(postprocessor_overrides or {})
    else:
        preprocessor_overrides, postprocessor_overrides = _legacy_groot_processor_overrides(
            config=config,
            dataset_stats=dataset_stats,
            preprocessor_overrides=preprocessor_overrides,
            postprocessor_overrides=postprocessor_overrides,
        )
    preprocessor = PolicyProcessorPipeline.from_pretrained(
        pretrained_model_name_or_path=pretrained_path,
        config_filename=preprocessor_config_filename,
        overrides=preprocessor_overrides,
        to_transition=batch_to_transition,
        to_output=transition_to_batch,
    )
    postprocessor = PolicyProcessorPipeline.from_pretrained(
        pretrained_model_name_or_path=pretrained_path,
        config_filename=postprocessor_config_filename,
        overrides=postprocessor_overrides,
        to_transition=policy_action_to_transition,
        to_output=transition_to_policy_action,
    )
    return preprocessor, postprocessor


def make_groot_pre_post_processors(
    config: GrootConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
) -> tuple[
    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
    PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
    """Create preprocessor and postprocessor for Groot policy.

    This creates a processing pipeline that transforms LeRobot data format into
    the format expected by Isaac-GR00T models:

    Preprocessing steps:
    1. Optional key renaming (dataset-specific key mapping)
    2. Add batch dimension to unbatched data
    3. Pack video/state/action/language/embodiment and apply optional min-max normalization before padding
    4. Encode video+language with Eagle VLM into intermediate eagle_content
    5. Collate eagle_content into batched eagle_* tensors
    6. Move tensors to device (GPU)

    NOTE: We optionally apply min-max normalization to STATE and ACTION using
    dataset-provided statistics prior to padding, mapping values to [-1, 1].
    This mirrors SO100-style preprocessing and keeps scales consistent with GR00T.

    Args:
        config: Groot configuration containing data_config, embodiment_tag, etc.
        dataset_stats: Optional per-key min/max statistics for normalization before padding.

    Returns:
        Tuple of (preprocessor, postprocessor) pipelines
    """

    if config.model_version == GROOT_N1_7:
        checkpoint_assets = _load_n1_7_checkpoint_processor_assets(config)
        action_horizon = (
            checkpoint_assets.max_action_horizon
            if checkpoint_assets is not None and checkpoint_assets.max_action_horizon is not None
            else min(config.chunk_size, 40)
        )
        valid_action_horizon = (
            checkpoint_assets.valid_action_horizon
            if checkpoint_assets is not None and checkpoint_assets.valid_action_horizon is not None
            else action_horizon
        )
        checkpoint_stats = checkpoint_assets.stats if checkpoint_assets is not None else None
        padded_stats = checkpoint_stats if _has_modality_stats(checkpoint_stats) else (dataset_stats or {})
        embodiment_mapping = (
            checkpoint_assets.embodiment_mapping
            if checkpoint_assets is not None
            else dict(N1_7_EMBODIMENT_MAPPING)
        )
        formalize_language = checkpoint_assets.formalize_language if checkpoint_assets is not None else True
        clip_outliers = checkpoint_assets.clip_outliers if checkpoint_assets is not None else True
        video_modality_keys = checkpoint_assets.video_modality_keys if checkpoint_assets is not None else None
        try:
            env_action_dim = int(config.output_features[ACTION].shape[0])
        except Exception:
            env_action_dim = 0
        state_cache_key = f"groot_n1_7:{config.embodiment_tag}"
        pack_step = GrootN17PackInputsStep(
            state_horizon=1,
            action_horizon=action_horizon,
            valid_action_horizon=valid_action_horizon,
            video_horizon=checkpoint_assets.video_horizon if checkpoint_assets is not None else None,
            max_state_dim=config.max_state_dim,
            max_action_dim=config.max_action_dim,
            language_key="task",
            formalize_language=formalize_language,
            embodiment_tag=config.embodiment_tag,
            embodiment_mapping=embodiment_mapping,
            normalize_min_max=True,
            stats=padded_stats,
            clip_outliers=clip_outliers,
            video_modality_keys=video_modality_keys,
            raw_stats=checkpoint_assets.raw_stats if checkpoint_assets is not None else None,
            modality_config=checkpoint_assets.modality_config if checkpoint_assets is not None else None,
            state_cache_key=state_cache_key,
        )

        input_steps: list[ProcessorStep] = [
            RenameObservationsProcessorStep(rename_map={}),
            AddBatchDimensionProcessorStep(),
            pack_step,
            GrootN17VLMEncodeStep(
                model_name=config.n1_7_backbone_model,
                image_crop_size=checkpoint_assets.image_crop_size if checkpoint_assets is not None else None,
                image_target_size=checkpoint_assets.image_target_size
                if checkpoint_assets is not None
                else None,
                shortest_image_edge=checkpoint_assets.shortest_image_edge
                if checkpoint_assets is not None
                else None,
                crop_fraction=checkpoint_assets.crop_fraction if checkpoint_assets is not None else None,
                use_albumentations=checkpoint_assets.use_albumentations
                if checkpoint_assets is not None
                else False,
            ),
            DeviceProcessorStep(device=config.device),
        ]
        if checkpoint_assets is None:
            action_decode_step: ProcessorStep = GrootActionUnpackUnnormalizeStep(
                env_action_dim=env_action_dim,
                stats=padded_stats,
                normalize_min_max=True,
                clip_normalized_action=True,
            )
        else:
            action_decode_step = GrootN17ActionDecodeStep(
                env_action_dim=env_action_dim,
                raw_stats=checkpoint_assets.raw_stats,
                modality_config=checkpoint_assets.modality_config,
                use_percentiles=checkpoint_assets.use_percentiles,
                use_relative_action=checkpoint_assets.use_relative_action,
                pack_step=pack_step,
                state_cache_key=state_cache_key,
                action_decode_transform=config.action_decode_transform,
            )

        output_steps: list[ProcessorStep] = [
            action_decode_step,
            DeviceProcessorStep(device="cpu"),
        ]
        return (
            PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
                steps=input_steps,
                name=POLICY_PREPROCESSOR_DEFAULT_NAME,
            ),
            PolicyProcessorPipeline[PolicyAction, PolicyAction](
                steps=output_steps,
                name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
                to_transition=policy_action_to_transition,
                to_output=transition_to_policy_action,
            ),
        )

    # Get horizon/dimension parameters from config
    # These should match the config used for the pretrained model
    # Default values match most GR00T configs (state_horizon=1, action_horizon=16)
    state_horizon = 1
    # CRITICAL: Pretrained GR00T models use action_horizon=16 max!
    # The model architecture hardcodes this limit
    action_horizon = min(config.chunk_size, 16)
    max_state_dim = config.max_state_dim
    max_action_dim = config.max_action_dim

    # Pass raw dataset_stats; normalization will occur inside pack step before padding
    padded_stats = dataset_stats or {}

    # Define feature specs for optional normalization steps
    _features: dict[str, PolicyFeature] = {
        # Observation features (only add those we may normalize)
        OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_horizon, max_state_dim)),
        # Action feature
        ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_horizon, max_action_dim)),
    }

    # Normalize STATE and ACTION with min_max (SO100-like default)
    _norm_map = {
        FeatureType.ACTION: NormalizationMode.MIN_MAX,
        FeatureType.STATE: NormalizationMode.MIN_MAX,
    }

    # Determine env action dimension from config (simple, object-like PolicyFeature)
    try:
        env_action_dim = int(config.output_features[ACTION].shape[0])
    except Exception:
        env_action_dim = 0

    input_steps: list[ProcessorStep] = [
        # 1. Rename keys if needed (e.g., dataset-specific camera names)
        # Leave empty for now - add mappings if your dataset uses different key names
        RenameObservationsProcessorStep(rename_map={}),
        # 2. Add batch dimension for single samples
        AddBatchDimensionProcessorStep(),
        # 3. Pack video/state/action/language/embodiment; apply optional min-max normalization before padding
        GrootPackInputsStep(
            state_horizon=state_horizon,
            action_horizon=action_horizon,
            max_state_dim=max_state_dim,
            max_action_dim=max_action_dim,
            language_key="task",
            formalize_language=False,
            embodiment_tag=config.embodiment_tag,
            normalize_min_max=True,
            stats=padded_stats,
        ),
        # 4. Eagle encode (creates eagle_content)
        GrootEagleEncodeStep(
            tokenizer_assets_repo=config.tokenizer_assets_repo,
        ),
        # 5. Collate eagle_content -> eagle_* tensors
        GrootEagleCollateStep(
            tokenizer_assets_repo=config.tokenizer_assets_repo,
        ),
        # 6. Move to device
        DeviceProcessorStep(device=config.device),
    ]

    # Postprocessing: slice to env action dim and unnormalize to env scale, then move to CPU
    output_steps: list[ProcessorStep] = [
        GrootActionUnpackUnnormalizeStep(
            env_action_dim=env_action_dim,
            stats=padded_stats,
            normalize_min_max=True,
        ),
        # Finally, move to CPU for env interaction
        DeviceProcessorStep(device="cpu"),
    ]

    return (
        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
            steps=input_steps,
            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
        ),
        PolicyProcessorPipeline[PolicyAction, PolicyAction](
            steps=output_steps,
            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
            to_transition=policy_action_to_transition,
            to_output=transition_to_policy_action,
        ),
    )


# GR00T specific processor steps


def _to_uint8_np_bthwc(img_t: torch.Tensor) -> np.ndarray:
    # img_t: (B, C, H, W) or (B, T, C, H, W), float in [0,1] or uint8
    if img_t.dtype.is_floating_point:
        img_t = (img_t.clamp(0, 1) * 255.0).to(torch.uint8)
    if img_t.dim() == 4:
        return rearrange(img_t.cpu().numpy(), "b c h w -> b 1 h w c")
    if img_t.dim() == 5:
        return rearrange(img_t.cpu().numpy(), "b t c h w -> b t h w c")
    raise ValueError(f"Expected image tensor shape (B, C, H, W) or (B, T, C, H, W), got {tuple(img_t.shape)}")


def _to_uint8_np_bhwc(img_t: torch.Tensor) -> np.ndarray:
    return _to_uint8_np_bthwc(img_t)[:, 0]


def _align_video_horizon(video: np.ndarray, horizon: int | None) -> np.ndarray:
    """Match the checkpoint video horizon by truncating or left-padding frames."""

    if horizon is None or horizon <= 0:
        return video
    current = video.shape[1]
    if current == horizon:
        return video
    if current > horizon:
        return video[:, -horizon:]
    pad = np.repeat(video[:, :1], horizon - current, axis=1)
    return np.concatenate([pad, video], axis=1)


def _infer_n1_7_batch_size_and_device(
    obs: dict[str, Any], action: torch.Tensor | None
) -> tuple[int, torch.device]:
    for value in list(obs.values()) + [action]:
        if isinstance(value, torch.Tensor):
            return value.shape[0], value.device
    video = obs.get("video")
    if isinstance(video, np.ndarray):
        return video.shape[0], torch.device("cpu")
    return 1, torch.device("cpu")


def _prepare_n1_7_language_batch(
    language: Any,
    batch_size: int,
    *,
    formalize_language: bool,
) -> list[str]:
    default_language = "Perform the task."
    if language is None or (isinstance(language, str) and language == ""):
        languages = [default_language] * batch_size
    elif isinstance(language, str):
        languages = [language] * batch_size
    elif isinstance(language, (list, tuple)):
        languages = list(language)
        if len(languages) == 0:
            languages = [default_language] * batch_size
        elif len(languages) == 1 and batch_size > 1:
            languages = languages * batch_size
        elif len(languages) != batch_size:
            raise ValueError(
                f"language batch has {len(languages)} entries, but GR00T N1.7 input batch has {batch_size}."
            )
    else:
        languages = [str(language)] * batch_size

    formatted = []
    for item in languages:
        text = str(item) if item else default_language
        if formalize_language:
            text = text.lower()
            text = "".join(ch for ch in text if ch.isalnum() or ch.isspace() or ch == "_")
        formatted.append(text)
    return formatted


def _build_eagle_processor(tokenizer_assets_repo: str = DEFAULT_TOKENIZER_ASSETS_REPO) -> ProcessorMixin:
    # Validate that the cache directory is ready. If not, instruct the user.
    cache_dir = HF_LEROBOT_HOME / tokenizer_assets_repo
    required = [
        cache_dir / "processor_config.json",
        cache_dir / "preprocessor_config.json",
        cache_dir / "image_processing_eagle2_5_vl_fast.py",
    ]
    if not all(p.exists() for p in required):
        raise FileNotFoundError(
            f"[GROOT] Eagle processor cache at '{cache_dir}' is not populated. "
            "Vendor files are copied during model creation. Create the policy/model first, "
            "or call ensure_eagle_cache_ready() before building processors."
        )
    proc = AutoProcessor.from_pretrained(
        str(cache_dir),
        trust_remote_code=True,
        fix_mistral_regex=False,
    )
    proc.tokenizer.padding_side = "left"
    return proc


def _build_n1_7_processor(model_name: str = GROOT_N1_7_BACKBONE_MODEL) -> ProcessorMixin:
    try:
        from transformers import (
            AutoTokenizer,
            Qwen2VLImageProcessorFast,
            Qwen3VLProcessor,
            Qwen3VLVideoProcessor,
        )
    except ImportError as exc:
        raise ImportError(
            "GR00T N1.7 preprocessing requires a transformers version with Qwen3-VL processor support. "
            "Install the GR00T optional dependencies with `pip install 'lerobot[groot]'`."
        ) from exc

    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    image_processor = Qwen2VLImageProcessorFast.from_pretrained(model_name, trust_remote_code=True)
    video_processor = Qwen3VLVideoProcessor.from_pretrained(model_name, trust_remote_code=True)
    proc = Qwen3VLProcessor(
        image_processor=image_processor,
        tokenizer=tokenizer,
        video_processor=video_processor,
        chat_template=tokenizer.chat_template,
    )
    proc.tokenizer.padding_side = "left"
    return proc


def _transform_n1_7_image_for_vlm(
    image: Image.Image,
    *,
    image_crop_size: list[int] | None,
    image_target_size: list[int] | None,
    shortest_image_edge: int | None,
    crop_fraction: float | None,
    use_albumentations: bool = False,
) -> Image.Image:
    if image_target_size is None:
        return image

    target_h, target_w = image_target_size
    if image.mode != "RGB":
        image = image.convert("RGB")

    if use_albumentations:
        try:
            import cv2
        except ImportError as exc:
            raise ImportError(
                "GR00T N1.7 checkpoints with use_albumentations=True require opencv-python-headless."
            ) from exc

        image_np = np.asarray(image)
        height, width = image_np.shape[:2]
        if height != width:
            square_edge = max(height, width)
            pad_h = square_edge - height
            pad_w = square_edge - width
            image_np = cv2.copyMakeBorder(
                image_np,
                pad_h // 2,
                pad_h - pad_h // 2,
                pad_w // 2,
                pad_w - pad_w // 2,
                cv2.BORDER_CONSTANT,
                value=(0, 0, 0),
            )

        resize_edge = shortest_image_edge or target_h
        if image_np.shape[:2] != (resize_edge, resize_edge):
            image_np = cv2.resize(image_np, (resize_edge, resize_edge), interpolation=cv2.INTER_AREA)

        if crop_fraction is None and image_crop_size is not None:
            crop_fraction = image_crop_size[0] / float(target_h)
        if crop_fraction is not None and 0.0 < crop_fraction < 1.0:
            height, width = image_np.shape[:2]
            crop_h = max(1, int(height * crop_fraction))
            crop_w = max(1, int(width * crop_fraction))
            top = max(0, (height - crop_h) // 2)
            left = max(0, (width - crop_w) // 2)
            image_np = image_np[top : top + crop_h, left : left + crop_w]

        if image_np.shape[:2] != (target_h, target_w):
            image_np = cv2.resize(image_np, (target_w, target_h), interpolation=cv2.INTER_AREA)
        return Image.fromarray(image_np)

    square_edge = max(image.width, image.height)
    if image.width != image.height:
        padded = Image.new("RGB", (square_edge, square_edge))
        left = (square_edge - image.width) // 2
        top = (square_edge - image.height) // 2
        padded.paste(image, (left, top))
        image = padded

    resize_edge = shortest_image_edge or target_h
    image = image.resize((resize_edge, resize_edge), Image.Resampling.BICUBIC)

    if crop_fraction is None and image_crop_size is not None:
        crop_fraction = image_crop_size[0] / float(target_h)
    if crop_fraction is not None and 0.0 < crop_fraction < 1.0:
        crop_w = max(1, int(round(image.width * crop_fraction)))
        crop_h = max(1, int(round(image.height * crop_fraction)))
        left = max(0, (image.width - crop_w) // 2)
        top = max(0, (image.height - crop_h) // 2)
        image = image.crop((left, top, left + crop_w, top + crop_h))

    if image.size != (target_w, target_h):
        image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
    return image


@dataclass
@ProcessorStepRegistry.register(name="groot_pack_inputs_v3")
class GrootPackInputsStep(ProcessorStep):
    state_horizon: int = 1
    action_horizon: int = 16
    max_state_dim: int = 64
    max_action_dim: int = 32
    language_key: str = "task"
    formalize_language: bool = False
    embodiment_tag: str = "new_embodiment"
    embodiment_mapping: dict[str, int] = field(
        default_factory=lambda: {
            "new_embodiment": 31,  # Match original GR00T EMBODIMENT_TAG_MAPPING
            "oxe_droid": 17,
            "agibot_genie1": 26,
            "gr1": 24,
            "so100": 2,
            "unitree_g1": 3,
        }
    )
    # Min-max normalization (SO100-like) applied BEFORE padding
    normalize_min_max: bool = True
    stats: dict[str, dict[str, Any]] | None = None

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        obs = transition.get(TransitionKey.OBSERVATION, {}) or {}
        comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}

        def _align_vec(vec: Any, target_dim: int, *, default: float) -> torch.Tensor:
            t = torch.as_tensor(vec)
            t = t.flatten().to(
                dtype=torch.float32,
                device=next(
                    (v.device for v in obs.values() if isinstance(v, torch.Tensor)), torch.device("cpu")
                ),
            )
            d = int(t.shape[-1]) if t.numel() > 0 else 0
            if d == target_dim:
                return t
            if d < target_dim:
                pad = torch.full((target_dim - d,), default, dtype=t.dtype, device=t.device)
                return torch.cat([t, pad], dim=0)
            return t[:target_dim]

        def _min_max_norm(x: torch.Tensor, key: str) -> torch.Tensor:
            if not self.normalize_min_max:
                return x
            if self.stats is None or key not in self.stats:
                return x
            stats_k = self.stats[key]
            last_dim = x.shape[-1]
            min_v = _align_vec(stats_k.get("min", torch.zeros(last_dim)), last_dim, default=0.0)
            max_v = _align_vec(stats_k.get("max", torch.ones(last_dim)), last_dim, default=1.0)
            denom = max_v - min_v
            mask = denom != 0
            safe_denom = torch.where(mask, denom, torch.ones_like(denom))
            mapped = 2 * (x - min_v) / safe_denom - 1
            return torch.where(mask, mapped, torch.zeros_like(mapped))

        # 1) Video (B, T=1, V, H, W, C) uint8
        img_keys = sorted([k for k in obs if k.startswith(OBS_IMAGES)])
        if not img_keys and OBS_IMAGE in obs:
            img_keys = [OBS_IMAGE]
        if img_keys:
            cams = [_to_uint8_np_bhwc(obs[k]) for k in img_keys]
            video = np.stack(cams, axis=1)  # (B, V, H, W, C)
            video = np.expand_dims(video, axis=1)  # (B, 1, V, H, W, C)
            # GR00T validates that video.shape[3] == 3 (channels), so reorder to (B, T, V, C, H, W)
            video = np.transpose(video, (0, 1, 2, 5, 3, 4))  # (B, 1, V, C, H, W)
            obs["video"] = video
            # Drop raw images to avoid confusion downstream
            for k in img_keys:
                obs.pop(k, None)

        # 2) Language (string)
        lang = comp.get(self.language_key)
        if isinstance(lang, list):
            lang = lang[0] if len(lang) > 0 else None
        if not lang:
            lang = "Perform the task."
        if self.formalize_language:
            lang = (lang or "").lower()
            lang = "".join(ch for ch in lang if ch.isalnum() or ch.isspace())
        comp["language"] = lang

        # 3) State/state_mask -> (B, 1, max_state_dim)
        if OBS_STATE in obs:
            state = obs[OBS_STATE]  # (B, D)
            if state.dim() != 2:
                raise ValueError(f"state must be (B, D), got {tuple(state.shape)}")
            bsz, d = state.shape
            # Normalize BEFORE padding
            if self.normalize_min_max:
                state = _min_max_norm(state, OBS_STATE)
            state = state.unsqueeze(1)  # (B, 1, D)
            if d > self.max_state_dim:
                state = state[:, :, : self.max_state_dim]
                d = self.max_state_dim
            elif d < self.max_state_dim:
                pad = torch.zeros(bsz, 1, self.max_state_dim - d, dtype=state.dtype, device=state.device)
                state = torch.cat([state, pad], dim=2)
            state_mask = torch.zeros(bsz, 1, self.max_state_dim, dtype=torch.bool, device=state.device)
            state_mask[:, :, :d] = True
            obs["state"] = state
            obs["state_mask"] = state_mask

        # 4) Action/action_mask -> (B, action_horizon, max_action_dim)
        action = transition.get(TransitionKey.ACTION)
        if isinstance(action, torch.Tensor):
            # Normalize BEFORE temporal expansion/padding
            if self.normalize_min_max:
                if action.dim() == 2:
                    action = _min_max_norm(action, ACTION)
                elif action.dim() == 3:
                    b, t, d = action.shape
                    flat = action.reshape(b * t, d)
                    flat = _min_max_norm(flat, ACTION)
                    action = flat.view(b, t, d)
            if action.dim() == 2:
                action = action.unsqueeze(1).repeat(1, self.action_horizon, 1)
            elif action.dim() == 3:
                b, t, d = action.shape
                if t < self.action_horizon:
                    last = action[:, -1:, :]
                    pad = last.repeat(1, self.action_horizon - t, 1)
                    action = torch.cat([action, pad], dim=1)
                elif t > self.action_horizon:
                    action = action[:, : self.action_horizon, :]
            else:
                raise ValueError(f"action must be (B, D) or (B, T, D), got {tuple(action.shape)}")

            b, t, d = action.shape
            if d > self.max_action_dim:
                action = action[:, :, : self.max_action_dim]
                d = self.max_action_dim
            elif d < self.max_action_dim:
                pad = torch.zeros(b, t, self.max_action_dim - d, dtype=action.dtype, device=action.device)
                action = torch.cat([action, pad], dim=2)
            action_mask = torch.zeros(b, t, self.max_action_dim, dtype=torch.bool, device=action.device)
            action_mask[:, :, :d] = True
            transition[TransitionKey.ACTION] = action
            comp["action_mask"] = action_mask

        # 5) Embodiment id as LongTensor (B,)
        emb_id = self.embodiment_mapping.get(self.embodiment_tag, 0)
        # Infer batch size/device from any tensor in obs or action
        bsz = None
        device = torch.device("cpu")
        for v in list(obs.values()) + [transition.get(TransitionKey.ACTION)]:
            if isinstance(v, torch.Tensor):
                bsz = v.shape[0]
                device = v.device
                break
        if bsz is None and "video" in obs and isinstance(obs["video"], np.ndarray):
            bsz = obs["video"].shape[0]
        if bsz is None:
            bsz = 1
        comp["embodiment_id"] = torch.full((bsz,), emb_id, dtype=torch.long, device=device)

        transition[TransitionKey.OBSERVATION] = obs
        transition[TransitionKey.COMPLEMENTARY_DATA] = comp
        return transition

    # Pipeline API requirement: declare how features change (we keep it simple)
    def transform_features(self, features):
        return features

    def get_config(self) -> dict[str, Any]:
        """
        Returns a serializable dictionary of the processor's configuration.

        Excludes 'stats' since they are saved separately via state_dict().
        """
        return {
            "state_horizon": self.state_horizon,
            "action_horizon": self.action_horizon,
            "max_state_dim": self.max_state_dim,
            "max_action_dim": self.max_action_dim,
            "language_key": self.language_key,
            "formalize_language": self.formalize_language,
            "embodiment_tag": self.embodiment_tag,
            "embodiment_mapping": self.embodiment_mapping,
            "normalize_min_max": self.normalize_min_max,
        }

    def state_dict(self) -> dict[str, torch.Tensor]:
        """
        Returns normalization statistics as a flat state dictionary.

        This enables saving stats to safetensors files, similar to normalizer_processor.
        """
        if not self.stats:
            return {}

        flat: dict[str, torch.Tensor] = {}
        for key, sub in self.stats.items():
            for stat_name, value in sub.items():
                tensor = torch.as_tensor(value).cpu()
                flat[f"{key}.{stat_name}"] = tensor
        return flat

    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
        """
        Loads normalization statistics from a flat state dictionary.

        This enables loading stats from safetensors files during from_pretrained.
        """
        if not state:
            return

        reconstructed: dict[str, dict[str, Any]] = {}
        for flat_key, tensor in state.items():
            if "." in flat_key:
                key, stat_name = flat_key.rsplit(".", 1)
                if key not in reconstructed:
                    reconstructed[key] = {}
                reconstructed[key][stat_name] = tensor

        if reconstructed:
            self.stats = reconstructed


@dataclass
@ProcessorStepRegistry.register(name="groot_n1_7_pack_inputs_v1")
class GrootN17PackInputsStep(ProcessorStep):
    """Pack LeRobot transitions into the raw tensor layout expected by N1.7.

    This step preserves the checkpoint's camera order, video horizon, language
    formatting, normalization statistics, action mask semantics, and embodiment
    id mapping before the Qwen3-VL processor sees the sample.
    """

    state_horizon: int = 1
    action_horizon: int = 40
    valid_action_horizon: int = 40
    video_horizon: int | None = None
    max_state_dim: int = 132
    max_action_dim: int = 132
    language_key: str = "task"
    formalize_language: bool = True
    embodiment_tag: str = "new_embodiment"
    embodiment_mapping: dict[str, int] = field(default_factory=lambda: dict(N1_7_EMBODIMENT_MAPPING))
    normalize_min_max: bool = True
    stats: dict[str, dict[str, Any]] | None = None
    clip_outliers: bool = True
    video_modality_keys: list[str] | None = None
    raw_stats: dict[str, Any] | None = None
    modality_config: dict[str, Any] | None = None
    state_cache_key: str = ""
    _last_raw_state: dict[str, np.ndarray] | None = field(default=None, init=False, repr=False)

    def _ordered_image_keys(self, obs: dict[str, Any]) -> list[str]:
        available = {key for key in obs if key.startswith(OBS_IMAGES)}
        if not available and OBS_IMAGE in obs:
            return [OBS_IMAGE]
        if not self.video_modality_keys:
            return sorted(available)

        ordered: list[str] = []
        for modality_key in self.video_modality_keys:
            candidates = [f"{OBS_IMAGES}.{modality_key}"]
            if modality_key == "wrist_image":
                candidates.append(f"{OBS_IMAGES}.image2")
            elif modality_key == "image":
                candidates.append(f"{OBS_IMAGES}.image")

            match = next((candidate for candidate in candidates if candidate in available), None)
            if match is not None:
                ordered.append(match)

        if not ordered:
            return sorted(available)
        return ordered

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        obs = transition.get(TransitionKey.OBSERVATION, {}) or {}
        comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}

        def _align_vec(vec: Any, target_dim: int, *, default: float) -> torch.Tensor:
            t = torch.as_tensor(vec)
            t = t.flatten().to(
                dtype=torch.float32,
                device=next(
                    (v.device for v in obs.values() if isinstance(v, torch.Tensor)), torch.device("cpu")
                ),
            )
            d = int(t.shape[-1]) if t.numel() > 0 else 0
            if d == target_dim:
                return t
            if d < target_dim:
                pad = torch.full((target_dim - d,), default, dtype=t.dtype, device=t.device)
                return torch.cat([t, pad], dim=0)
            return t[:target_dim]

        def _min_max_norm(x: torch.Tensor, key: str) -> torch.Tensor:
            if not self.normalize_min_max or self.stats is None or key not in self.stats:
                return x
            stats_k = self.stats[key]
            last_dim = x.shape[-1]
            min_v = _align_vec(stats_k.get("min", torch.zeros(last_dim)), last_dim, default=0.0)
            max_v = _align_vec(stats_k.get("max", torch.ones(last_dim)), last_dim, default=1.0)
            denom = max_v - min_v
            mask = denom != 0
            safe_denom = torch.where(mask, denom, torch.ones_like(denom))
            mapped = 2 * (x - min_v) / safe_denom - 1
            normalized = torch.where(mask, mapped, torch.zeros_like(mapped))
            if self.clip_outliers:
                normalized = normalized.clamp(-1.0, 1.0)
            return normalized

        def _cache_raw_state(state: torch.Tensor) -> None:
            if self.modality_config is None or self.raw_stats is None:
                return
            state_config = self.modality_config.get("state", {})
            if not isinstance(state_config, dict):
                return
            state_keys = state_config.get("modality_keys", [])
            if not isinstance(state_keys, list):
                return

            raw_state = state.detach().cpu().float().numpy()
            start_idx = 0
            grouped: dict[str, np.ndarray] = {}
            for key in state_keys:
                if not isinstance(key, str):
                    continue
                key_stats = self.raw_stats.get("state", {}).get(key, {})
                dim = len(key_stats.get("mean") or key_stats.get("min") or key_stats.get("q01") or [])
                if dim <= 0:
                    continue
                grouped[key] = raw_state[:, start_idx : start_idx + dim]
                start_idx += dim
            if grouped:
                self._last_raw_state = grouped
                _N1_7_RAW_STATE_CACHE[_n1_7_state_cache_key(self.state_cache_key)] = grouped

        img_keys = self._ordered_image_keys(obs)
        if img_keys:
            cams = [_align_video_horizon(_to_uint8_np_bthwc(obs[k]), self.video_horizon) for k in img_keys]
            video = np.stack(cams, axis=2)  # (B, T, V, H, W, C)
            obs["video"] = video
            image_keys_to_remove = [key for key in obs if key.startswith(OBS_IMAGES)]
            if OBS_IMAGE in obs:
                image_keys_to_remove.append(OBS_IMAGE)
            for k in image_keys_to_remove:
                obs.pop(k, None)

        bsz, _device = _infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION))
        comp["language"] = _prepare_n1_7_language_batch(
            comp.get(self.language_key),
            bsz,
            formalize_language=self.formalize_language,
        )

        if OBS_STATE in obs:
            state = obs[OBS_STATE]
            if state.dim() != 2:
                raise ValueError(f"state must be (B, D), got {tuple(state.shape)}")
            bsz, dim = state.shape
            if dim > self.max_state_dim:
                raise ValueError(f"State dimension {dim} exceeds max_state_dim {self.max_state_dim}.")
            _cache_raw_state(state)
            if self.normalize_min_max:
                state = _min_max_norm(state, OBS_STATE)
            state = state.unsqueeze(1)
            if dim < self.max_state_dim:
                pad = torch.zeros(bsz, 1, self.max_state_dim - dim, dtype=state.dtype, device=state.device)
                state = torch.cat([state, pad], dim=2)
            obs["state"] = state

        action = transition.get(TransitionKey.ACTION)
        if isinstance(action, torch.Tensor):
            if action.dim() == 2:
                action = action.unsqueeze(1)
            elif action.dim() == 3:
                pass
            else:
                raise ValueError(f"action must be (B, D) or (B, T, D), got {tuple(action.shape)}")

            bsz, horizon, dim = action.shape
            if horizon > self.action_horizon:
                raise ValueError(f"Action horizon {horizon} exceeds action_horizon {self.action_horizon}.")
            if dim > self.max_action_dim:
                raise ValueError(f"Action dimension {dim} exceeds max_action_dim {self.max_action_dim}.")
            if self.normalize_min_max:
                flat = _min_max_norm(action.reshape(bsz * horizon, dim), ACTION)
                action = flat.view(bsz, horizon, dim)
            valid_dim = min(dim, self.max_action_dim)
            valid_horizon = min(horizon, self.valid_action_horizon, self.action_horizon)
            if dim < self.max_action_dim:
                pad = torch.zeros(
                    bsz, horizon, self.max_action_dim - dim, dtype=action.dtype, device=action.device
                )
                action = torch.cat([action, pad], dim=2)
            if horizon < self.action_horizon:
                pad = torch.zeros(
                    bsz,
                    self.action_horizon - horizon,
                    self.max_action_dim,
                    dtype=action.dtype,
                    device=action.device,
                )
                action = torch.cat([action, pad], dim=1)
                horizon = self.action_horizon
            if valid_horizon < horizon:
                action = action.clone()
                action[:, valid_horizon:, :] = 0
            action_mask = torch.zeros(
                bsz, horizon, self.max_action_dim, dtype=torch.float32, device=action.device
            )
            action_mask[:, :valid_horizon, :valid_dim] = 1.0
            transition[TransitionKey.ACTION] = action
            comp["action_mask"] = action_mask

        emb_id = self.embodiment_mapping.get(self.embodiment_tag, 0)
        bsz, device = _infer_n1_7_batch_size_and_device(obs, transition.get(TransitionKey.ACTION))
        if "action_mask" not in comp:
            action_mask = torch.zeros(bsz, self.action_horizon, dtype=torch.float32, device=device)
            valid_horizon = min(self.valid_action_horizon, self.action_horizon)
            action_mask[:, :valid_horizon] = 1.0
            comp["action_mask"] = action_mask
        comp["embodiment_id"] = torch.full((bsz,), emb_id, dtype=torch.int32, device=device)

        transition[TransitionKey.OBSERVATION] = obs
        transition[TransitionKey.COMPLEMENTARY_DATA] = comp
        return transition

    def transform_features(self, features):
        return features

    def get_config(self) -> dict[str, Any]:
        return {
            "state_horizon": self.state_horizon,
            "action_horizon": self.action_horizon,
            "valid_action_horizon": self.valid_action_horizon,
            "video_horizon": self.video_horizon,
            "max_state_dim": self.max_state_dim,
            "max_action_dim": self.max_action_dim,
            "language_key": self.language_key,
            "formalize_language": self.formalize_language,
            "embodiment_tag": self.embodiment_tag,
            "embodiment_mapping": self.embodiment_mapping,
            "normalize_min_max": self.normalize_min_max,
            "clip_outliers": self.clip_outliers,
            "video_modality_keys": self.video_modality_keys,
            "raw_stats": self.raw_stats,
            "modality_config": self.modality_config,
            "state_cache_key": self.state_cache_key,
        }

    def get_cached_raw_state(self) -> dict[str, np.ndarray] | None:
        """Return the latest unnormalized state split by checkpoint modality key."""

        return self._last_raw_state

    def state_dict(self) -> dict[str, torch.Tensor]:
        if not self.stats:
            return {}

        flat: dict[str, torch.Tensor] = {}
        for key, sub in self.stats.items():
            for stat_name, value in sub.items():
                flat[f"{key}.{stat_name}"] = torch.as_tensor(value).cpu()
        return flat

    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
        if not state:
            return
        reconstructed: dict[str, dict[str, Any]] = {}
        for flat_key, tensor in state.items():
            if "." in flat_key:
                key, stat_name = flat_key.rsplit(".", 1)
                reconstructed.setdefault(key, {})[stat_name] = tensor
        if reconstructed:
            self.stats = reconstructed


@dataclass
@ProcessorStepRegistry.register(name="groot_n1_7_vlm_encode_v1")
class GrootN17VLMEncodeStep(ProcessorStep):
    """Tokenize N1.7's packed video-language prompt with the Qwen3-VL processor.

    The packed video has shape ``(B, T, V, H, W, C)``. Each frame/view becomes
    an image item in the same chat message so the resulting image tokens match
    the temporal VLM packing used by Isaac-GR00T.
    """

    model_name: str = GROOT_N1_7_BACKBONE_MODEL
    image_crop_size: list[int] | None = None
    image_target_size: list[int] | None = None
    shortest_image_edge: int | None = None
    crop_fraction: float | None = None
    use_albumentations: bool = False
    _proc: ProcessorMixin | None = field(default=None, init=False, repr=False)

    @property
    def proc(self) -> ProcessorMixin:
        if self._proc is None:
            self._proc = _build_n1_7_processor(self.model_name)
        return self._proc

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        obs = transition.get(TransitionKey.OBSERVATION, {}) or {}
        comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}
        video = obs.get("video")
        if video is None:
            return transition

        languages = _prepare_n1_7_language_batch(
            comp.get("language"),
            video.shape[0],
            formalize_language=False,
        )

        texts: list[str] = []
        images: list[Image.Image] = []
        for batch_idx in range(video.shape[0]):
            sample = video[batch_idx]  # (T, V, H, W, C)
            sample_images = [
                _transform_n1_7_image_for_vlm(
                    Image.fromarray(sample[timestep, view_idx]),
                    image_crop_size=self.image_crop_size,
                    image_target_size=self.image_target_size,
                    shortest_image_edge=self.shortest_image_edge,
                    crop_fraction=self.crop_fraction,
                    use_albumentations=self.use_albumentations,
                )
                for timestep in range(sample.shape[0])
                for view_idx in range(sample.shape[1])
            ]
            conversation = [
                {
                    "role": "user",
                    "content": [
                        *[{"type": "image", "image": image} for image in sample_images],
                        {"type": "text", "text": languages[batch_idx]},
                    ],
                }
            ]
            texts.append(
                self.proc.apply_chat_template(
                    conversation,
                    tokenize=False,
                    add_generation_prompt=False,
                )
            )
            images.extend(sample_images)

        encoded = self.proc(text=texts, images=images, return_tensors="pt", padding=True)
        for key, value in encoded.items():
            comp[key] = value
        obs.pop("video", None)
        transition[TransitionKey.OBSERVATION] = obs
        transition[TransitionKey.COMPLEMENTARY_DATA] = comp
        return transition

    def transform_features(self, features):
        return features

    def get_config(self) -> dict[str, Any]:
        return {
            "model_name": self.model_name,
            "image_crop_size": self.image_crop_size,
            "image_target_size": self.image_target_size,
            "shortest_image_edge": self.shortest_image_edge,
            "crop_fraction": self.crop_fraction,
            "use_albumentations": self.use_albumentations,
        }


@dataclass
@ProcessorStepRegistry.register(name="groot_eagle_encode_v3")
class GrootEagleEncodeStep(ProcessorStep):
    tokenizer_assets_repo: str = DEFAULT_TOKENIZER_ASSETS_REPO
    _proc: ProcessorMixin | None = field(default=None, init=False, repr=False)

    @property
    def proc(self) -> ProcessorMixin:
        if self._proc is None:
            self._proc = _build_eagle_processor(self.tokenizer_assets_repo)
        return self._proc

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        obs = transition.get(TransitionKey.OBSERVATION, {}) or {}
        comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}

        if "video" not in obs:
            return transition

        video = obs["video"]  # (B, T, V, H, W, C) uint8
        lang = comp.get("language", "Perform the task.")
        if isinstance(lang, list):
            lang = lang[0] if len(lang) > 0 else "Perform the task."

        bsz = video.shape[0]
        eagle_contents: list[dict[str, Any]] = []
        for b in range(bsz):
            vt = video[b]  # (T, V, C, H, W) after reorder
            if vt.ndim != 5:
                # Fallback: assume (T, V, H, W, C)
                t, v, h, w, c = vt.shape
                flat = rearrange(vt, "t v h w c -> (t v) h w c")
            else:
                t, v, c, h, w = vt.shape
                flat = rearrange(vt, "t v c h w -> (t v) h w c")
            images = [Image.fromarray(flat[i]) for i in range(t * v)]
            # Format language as string list representation to match Original GROOT
            lang_formatted = str([lang])
            text_content = [{"type": "text", "text": lang_formatted}]
            image_content = [{"type": "image", "image": img} for img in images]
            conv = [{"role": "user", "content": image_content + text_content}]
            text_list = [self.proc.apply_chat_template(conv, tokenize=False, add_generation_prompt=True)]
            img_inputs, vid_inputs = self.proc.process_vision_info(conv)
            eagle_contents.append(
                {
                    "text_list": text_list,
                    "image_inputs": img_inputs,
                    "video_inputs": vid_inputs,
                }
            )

        comp["eagle_content"] = eagle_contents
        transition[TransitionKey.OBSERVATION] = obs
        transition[TransitionKey.COMPLEMENTARY_DATA] = comp
        return transition

    # Pipeline API requirement: declare how features change (no schema change here)
    def transform_features(self, features):
        return features


# Original GR00T-style collate: converts eagle_content -> eagle_* tensors
def collate(features: list[dict[str, Any]], eagle_processor: ProcessorMixin) -> dict[str, Any]:
    batch: dict[str, Any] = {}
    keys = features[0].keys()

    for key in keys:
        values = [elem[key] for elem in features]

        if key == "eagle_content":
            text_list: list[str] = []
            image_inputs: list[Any] = []
            for v in values:
                curr_text_list = v["text_list"]
                curr_image_inputs = v["image_inputs"]
                text_list += curr_text_list
                image_inputs += curr_image_inputs
            eagle_inputs = eagle_processor(
                text=text_list,
                images=image_inputs,
                images_kwargs={"min_dynamic_tiles": 1, "max_dynamic_tiles": 1, "use_thumbnail": False},
                return_tensors="pt",
                padding=True,
            )
            for k, v in eagle_inputs.items():
                k = "eagle_" + k
                batch[k] = v
        elif key in ("pixel_values", "image_grid_thw", "attention_mask", "input_ids"):
            # Concat in existing batch dimension.
            batch[key] = torch.cat(values)
        else:
            # state, state_mask, action and action_mask.
            # Stack to form the batch dimension.
            batch[key] = torch.from_numpy(np.stack(values))
    return batch


@dataclass
@ProcessorStepRegistry.register(name="groot_eagle_collate_v3")
class GrootEagleCollateStep(ProcessorStep):
    tokenizer_assets_repo: str = DEFAULT_TOKENIZER_ASSETS_REPO
    _proc: ProcessorMixin | None = field(default=None, init=False, repr=False)

    @property
    def proc(self) -> ProcessorMixin:
        if self._proc is None:
            self._proc = _build_eagle_processor(self.tokenizer_assets_repo)
        return self._proc

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        obs = transition.get(TransitionKey.OBSERVATION, {}) or {}
        comp = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}
        contents = comp.get("eagle_content")
        if not contents:
            return transition

        # Build features list as original API expects: one dict per batch item
        features = [{"eagle_content": content} for content in contents]
        batched = collate(features, self.proc)

        # Inject eagle_* tensors and remove the temporary content and raw video to free memory
        for k, v in batched.items():
            comp[k] = v
        comp.pop("eagle_content", None)
        obs.pop(
            "video", None
        )  # The video has been fully encoded into eagle_* tensors, so we don't need the raw video anymore
        transition[TransitionKey.OBSERVATION] = obs
        transition[TransitionKey.COMPLEMENTARY_DATA] = comp
        return transition

    def transform_features(self, features):
        return features


def _stat_dim_from_entry(entry: dict[str, Any]) -> int:
    for stat_name in ("mean", "q01", "min", "max", "std"):
        value = entry.get(stat_name)
        if isinstance(value, list) and len(value) > 0:
            return len(value)
    return 0


def _n1_7_decode_stats_for_action(
    raw_stats: dict[str, Any],
    key: str,
    action_config: dict[str, Any],
    *,
    use_relative_action: bool,
    use_percentiles: bool,
) -> tuple[np.ndarray, np.ndarray]:
    """Select the min/max arrays needed to decode one checkpoint action group."""

    is_relative = use_relative_action and _config_value(action_config.get("rep")) == "relative"
    modality = "relative_action" if is_relative else "action"
    stats = raw_stats.get(modality, {}).get(key, {})
    if not isinstance(stats, dict):
        raise KeyError(f"Missing N1.7 statistics for {modality}.{key}")
    min_name = "min" if is_relative else ("q01" if use_percentiles else "min")
    max_name = "max" if is_relative else ("q99" if use_percentiles else "max")
    if min_name not in stats or max_name not in stats:
        raise KeyError(f"Missing '{min_name}'/'{max_name}' statistics for {modality}.{key}")
    return np.asarray(stats[min_name], dtype=np.float32), np.asarray(stats[max_name], dtype=np.float32)


def _unnormalize_min_max(action: np.ndarray, min_v: np.ndarray, max_v: np.ndarray) -> np.ndarray:
    return (np.clip(action, -1.0, 1.0) + 1.0) * 0.5 * (max_v - min_v) + min_v


def _n1_7_decode_valid_horizon(action_config: dict[str, Any], action_np: np.ndarray) -> int | None:
    if action_np.ndim != 3:
        return None
    delta_indices = action_config.get("delta_indices", [])
    if not isinstance(delta_indices, list) or not delta_indices:
        return None
    return max(1, min(action_np.shape[1], len(delta_indices)))


def _rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray:
    rows = rot6d.reshape(2, 3).astype(np.float64)
    row1 = rows[0] / np.linalg.norm(rows[0])
    row2 = rows[1] - np.dot(row1, rows[1]) * row1
    row2 = row2 / np.linalg.norm(row2)
    row3 = np.cross(row1, row2)
    return np.vstack([row1, row2, row3])


def _xyz_rot6d_to_homogeneous(xyz_rot6d: np.ndarray) -> np.ndarray:
    transform = np.eye(4, dtype=np.float64)
    transform[:3, :3] = _rot6d_to_matrix(xyz_rot6d[3:])
    transform[:3, 3] = xyz_rot6d[:3]
    return transform


def _homogeneous_to_xyz_rot6d(transform: np.ndarray) -> np.ndarray:
    return np.concatenate([transform[:3, 3], transform[:2, :3].reshape(-1)], axis=0)


def _relative_eef_to_absolute(action: np.ndarray, reference_state: np.ndarray) -> np.ndarray:
    """Convert relative EEF deltas in xyz+rot6d format to absolute EEF poses."""

    out = np.empty_like(action, dtype=np.float64)
    for batch_idx in range(action.shape[0]):
        reference = _xyz_rot6d_to_homogeneous(reference_state[batch_idx])
        for timestep in range(action.shape[1]):
            relative = _xyz_rot6d_to_homogeneous(action[batch_idx, timestep])
            out[batch_idx, timestep] = _homogeneous_to_xyz_rot6d(reference @ relative)
    return out.astype(np.float32)


def _n1_7_action_group_slice(
    action_keys: list[Any], decoded_groups: dict[str, np.ndarray], target_key: str
) -> slice:
    start_idx = 0
    for key in action_keys:
        if not isinstance(key, str) or key not in decoded_groups:
            continue
        dim = decoded_groups[key].shape[-1]
        end_idx = start_idx + dim
        if key == target_key:
            return slice(start_idx, end_idx)
        start_idx = end_idx

    raise KeyError(f"Missing N1.7 action group '{target_key}' required by action decode transform.")


def _apply_n1_7_action_decode_transform(
    decoded: np.ndarray,
    *,
    transform: str | None,
    action_keys: list[Any],
    decoded_groups: dict[str, np.ndarray],
) -> np.ndarray:
    if transform is None:
        return decoded

    if transform == GROOT_ACTION_DECODE_TRANSFORM_LIBERO:
        gripper_slice = _n1_7_action_group_slice(action_keys, decoded_groups, "gripper")
        if gripper_slice.stop is None or gripper_slice.stop > decoded.shape[-1]:
            raise ValueError(
                "N1.7 LIBERO action decode transform requested, but the decoded gripper action "
                "is outside the sliced environment action."
            )
        if gripper_slice.stop - gripper_slice.start != 1:
            raise ValueError("N1.7 LIBERO action decode transform expects a scalar gripper action.")

        transformed = decoded.copy()
        gripper = transformed[..., gripper_slice]
        transformed[..., gripper_slice] = -np.sign(2.0 * gripper - 1.0)
        return transformed

    raise ValueError(f"Unsupported N1.7 action decode transform '{transform}'.")


@dataclass
@ProcessorStepRegistry.register(name="groot_n1_7_action_decode_v1")
class GrootN17ActionDecodeStep(ProcessorStep):
    """Decode the full 132-D N1.7 model action back to environment actions.

    N1.7 predicts checkpoint-order action groups. This step unnormalizes each
    group with the checkpoint stats, converts relative groups to absolute values
    using the raw state cached during packing, concatenates groups in checkpoint
    order, and finally slices to the environment action dimension.
    """

    env_action_dim: int = 0
    raw_stats: dict[str, Any] | None = None
    modality_config: dict[str, Any] | None = None
    use_percentiles: bool = False
    use_relative_action: bool = False
    state_cache_key: str = ""
    action_decode_transform: str | None = None
    pack_step: GrootN17PackInputsStep | None = field(default=None, repr=False)

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        action = transition.get(TransitionKey.ACTION)
        if not isinstance(action, torch.Tensor):
            return transition
        if self.raw_stats is None or self.modality_config is None:
            return transition

        action_config = self.modality_config.get("action", {})
        if not isinstance(action_config, dict):
            return transition
        action_keys = action_config.get("modality_keys", [])
        action_configs = action_config.get("action_configs", [])
        if not isinstance(action_keys, list) or not isinstance(action_configs, list):
            return transition

        action_np = action.detach().cpu().float().numpy()
        valid_horizon = _n1_7_decode_valid_horizon(action_config, action_np)
        if valid_horizon is not None:
            action_np = action_np[:, :valid_horizon]
        decoded_groups: dict[str, np.ndarray] = {}
        start_idx = 0
        for idx, key in enumerate(action_keys):
            if not isinstance(key, str):
                continue
            stats_entry = self.raw_stats.get("action", {}).get(key, {})
            if not isinstance(stats_entry, dict):
                continue
            dim = _stat_dim_from_entry(stats_entry)
            if dim <= 0:
                continue
            cfg = (
                action_configs[idx]
                if idx < len(action_configs) and isinstance(action_configs[idx], dict)
                else {}
            )
            normalized = action_np[..., start_idx : start_idx + dim]
            min_v, max_v = _n1_7_decode_stats_for_action(
                self.raw_stats,
                key,
                cfg,
                use_relative_action=self.use_relative_action,
                use_percentiles=self.use_percentiles,
            )
            decoded_groups[key] = _unnormalize_min_max(normalized, min_v, max_v)
            start_idx += dim

        if self.use_relative_action:
            raw_state = self.pack_step.get_cached_raw_state() if self.pack_step is not None else None
            if raw_state is None:
                raw_state = _N1_7_RAW_STATE_CACHE.get(_n1_7_state_cache_key(self.state_cache_key))
            if raw_state is None:
                raise RuntimeError(
                    "GrootN17ActionDecodeStep requires cached raw state from GrootN17PackInputsStep "
                    "to convert relative N1.7 actions back to absolute actions."
                )
            for idx, key in enumerate(action_keys):
                if not isinstance(key, str) or key not in decoded_groups or idx >= len(action_configs):
                    continue
                cfg = action_configs[idx]
                if not isinstance(cfg, dict) or _config_value(cfg.get("rep")) != "relative":
                    continue
                state_key = cfg.get("state_key") or key
                if state_key not in raw_state:
                    raise KeyError(f"Missing cached raw state '{state_key}' for relative N1.7 action '{key}'")
                reference = raw_state[state_key]
                action_type = _config_value(cfg.get("type"))
                action_format = _config_value(cfg.get("format"))
                if action_type == "non_eef":
                    decoded_groups[key] = decoded_groups[key] + reference[:, None, :]
                elif action_type == "eef" and action_format == "xyz+rot6d":
                    decoded_groups[key] = _relative_eef_to_absolute(decoded_groups[key], reference)
                else:
                    raise ValueError(f"Unsupported relative N1.7 action config for '{key}': {cfg}")

        if not decoded_groups:
            return transition

        decoded = np.concatenate(
            [decoded_groups[key] for key in action_keys if isinstance(key, str) and key in decoded_groups],
            axis=-1,
        )
        if self.env_action_dim and decoded.shape[-1] > self.env_action_dim:
            decoded = decoded[..., : self.env_action_dim]
        decoded = _apply_n1_7_action_decode_transform(
            decoded,
            transform=self.action_decode_transform,
            action_keys=action_keys,
            decoded_groups=decoded_groups,
        )
        new_transition = transition.copy()
        new_transition[TransitionKey.ACTION] = torch.as_tensor(
            decoded, dtype=action.dtype, device=action.device
        )
        return new_transition

    def transform_features(self, features):
        return features

    def get_config(self) -> dict[str, Any]:
        return {
            "env_action_dim": self.env_action_dim,
            "raw_stats": self.raw_stats,
            "modality_config": self.modality_config,
            "use_percentiles": self.use_percentiles,
            "use_relative_action": self.use_relative_action,
            "state_cache_key": self.state_cache_key,
            "action_decode_transform": self.action_decode_transform,
        }


@dataclass
@ProcessorStepRegistry.register(name="groot_action_unpack_unnormalize_v1")
class GrootActionUnpackUnnormalizeStep(ProcessorStep):
    env_action_dim: int = 0
    # Apply inverse of min-max normalization if it was used in preprocessor
    normalize_min_max: bool = True
    stats: dict[str, dict[str, Any]] | None = None
    clip_normalized_action: bool = False
    libero_gripper_action: bool = False
    libero_gripper_binarize: bool = True

    def __call__(self, transition: EnvTransition) -> EnvTransition:
        # Expect model outputs to be in TransitionKey.ACTION as (B, T, D_model)
        action = transition.get(TransitionKey.ACTION)
        if not isinstance(action, torch.Tensor):
            return transition

        # Slice to env dimension while preserving an optional action horizon.
        # Sync rollout postprocesses selected actions as (B, D); RTC postprocesses
        # chunks as (B, T, D), matching Isaac-GR00T's decode_action contract.
        if self.env_action_dim and action.shape[-1] >= self.env_action_dim:
            action = action[..., : self.env_action_dim]

        # Inverse min-max normalization mirroring _min_max_norm:
        # forward: y = 2 * (x - min) / denom - 1, with y=0 when denom==0
        # inverse: x = (y+1)/2 * denom + min, and when denom==0 -> x = min
        if self.normalize_min_max and self.stats is not None:
            if self.clip_normalized_action:
                action = action.clamp(-1.0, 1.0)
            stats_k = self.stats.get(ACTION, {})
            d = action.shape[-1]
            min_v = torch.as_tensor(
                stats_k.get("min", torch.zeros(d)), dtype=action.dtype, device=action.device
            )
            max_v = torch.as_tensor(
                stats_k.get("max", torch.ones(d)), dtype=action.dtype, device=action.device
            )
            if min_v.numel() != d:
                min_v = torch.nn.functional.pad(min_v.flatten()[:d], (0, max(0, d - min_v.numel())))
                min_v = min_v.to(action.device, dtype=action.dtype)
            if max_v.numel() != d:
                max_v = torch.nn.functional.pad(max_v.flatten()[:d], (0, max(0, d - max_v.numel())))
                max_v = max_v.to(action.device, dtype=action.dtype)
            denom = max_v - min_v
            mask = denom != 0
            safe_denom = torch.where(mask, denom, torch.ones_like(denom))
            inv = (action + 1.0) * 0.5 * safe_denom + min_v
            action = torch.where(mask, inv, min_v)

        if self.libero_gripper_action and action.shape[-1] >= 7:
            gripper = action[..., -1]
            if self.libero_gripper_binarize:
                gripper = -torch.sign(2.0 * gripper - 1.0)
            else:
                gripper = -(2.0 * gripper - 1.0)
            action = action.clone()
            action[..., -1] = gripper

        transition[TransitionKey.ACTION] = action
        return transition

    def transform_features(self, features):
        return features

    def get_config(self) -> dict[str, Any]:
        """
        Returns a serializable dictionary of the processor's configuration.

        Excludes 'stats' since they are saved separately via state_dict().
        """
        return {
            "env_action_dim": self.env_action_dim,
            "normalize_min_max": self.normalize_min_max,
            "clip_normalized_action": self.clip_normalized_action,
            "libero_gripper_action": self.libero_gripper_action,
            "libero_gripper_binarize": self.libero_gripper_binarize,
        }

    def state_dict(self) -> dict[str, torch.Tensor]:
        """
        Returns normalization statistics as a flat state dictionary.

        This enables saving stats to safetensors files, similar to normalizer_processor.
        """
        if not self.stats:
            return {}

        flat: dict[str, torch.Tensor] = {}
        for key, sub in self.stats.items():
            for stat_name, value in sub.items():
                tensor = torch.as_tensor(value).cpu()
                flat[f"{key}.{stat_name}"] = tensor
        return flat

    def load_state_dict(self, state: dict[str, torch.Tensor]) -> None:
        """
        Loads normalization statistics from a flat state dictionary.

        This enables loading stats from safetensors files during from_pretrained.
        """
        if not state:
            return

        reconstructed: dict[str, dict[str, Any]] = {}
        for flat_key, tensor in state.items():
            if "." in flat_key:
                key, stat_name = flat_key.rsplit(".", 1)
                if key not in reconstructed:
                    reconstructed[key] = {}
                reconstructed[key][stat_name] = tensor

        if reconstructed:
            self.stats = reconstructed
