scripts/compute_norm_stats_fast.py

"""Compute normalization statistics for a config.

This script is used to compute the normalization statistics for a given config. It
will compute the mean and standard deviation of the data in the dataset and save it
to the config assets directory.
"""
import dataclasses
import pathlib
import numpy as np
import tqdm
import tyro

import openpi_value.models.model as _model
import openpi_value.shared.normalize as normalize
import openpi_value.training.config as _config
import openpi_value.training.data_loader as _data_loader
import openpi_value.transforms as transforms

# * python scripts/compute_norm_stats_fast.py --config-name Compute_norm
# * Remember to modify the Compute_norm beforehand

@dataclasses.dataclass(frozen=True)
class FakeInputs(transforms.DataTransformFn):
    """Inputs for the CustomAgilex policy.

    Expected inputs:
    - images: dict[name, img] where img is [channel, height, width]. name must be in EXPECTED_CAMERAS.
    - state: [14]
    - actions: [action_horizon, 14]
    """

    # The action dimension of the model. Will be used to pad state and actions.
    action_dim: int

    # Determines which model will be used.
    model_type: _model.ModelType = _model.ModelType.PI0
    
    # if set all state to zeros
    mask_state: bool = False

    # if convert to eef position
    convert_to_eef_position: bool = False



    def __call__(self, data: dict) -> dict:
        # We only mask padding for pi0 model, not pi0-FAST
        mask_padding = self.model_type == _model.ModelType.PI0

        # Pad the proprioceptive input to the action dimension of the model
        state = transforms.pad_to_dim(data["state"], self.action_dim)
        # Ensure state has correct shape [batch_size, state_dim]
        state = state.squeeze()


        # * We need to mask out extremely large values.
        state = np.where(state > np.pi, 0, state)
        state = np.where(state < -np.pi, 0, state)

        # Prepare inputs dictionary
        masked_state = np.zeros_like(state) if self.mask_state else state
        inputs = {
            "state": masked_state,
        }

        # Add actions if present
        if "actions" in data:
            actions = transforms.pad_to_dim(data["actions"], self.action_dim)
            
            # * We need to mask out extremely large values.
            actions = np.where(actions > np.pi, 0, actions)
            actions = np.where(actions < -np.pi, 0, actions)
            
            if mask_padding:
                # Create action mask for padding
                action_mask = np.ones_like(actions, dtype=bool)
                action_mask[:, self.action_dim:] = False
                inputs["action_mask"] = action_mask

            inputs["actions"] = actions.squeeze()

        # Add prompt if present
        if "prompt" in data:
            inputs["prompt"] = data["prompt"]

        return inputs


@dataclasses.dataclass(frozen=True)
class CustomAgilexOutputs(transforms.DataTransformFn):
    """Outputs for the CustomAgilex policy."""

    def __call__(self, data: dict) -> dict:
        # Return the first 14 dimensions of actions (13 joints + 1 gripper)
        return {"actions": np.asarray(data["actions"][:, :14])} 


class RemoveStrings(transforms.DataTransformFn):
    def __call__(self, x: dict) -> dict:
        return {k: v for k, v in x.items() if not np.issubdtype(np.asarray(v).dtype, np.str_)}


def create_torch_dataloader(
    data_config: _config.DataConfig,
    action_horizon: int,
    batch_size: int,
    model_config: _model.BaseModelConfig,
    num_workers: int,
    max_frames: int | None = None,
) -> tuple[_data_loader.Dataset, int]:
    if data_config.repo_id is None:
        raise ValueError("Data config must have a repo_id")
    
    dataset = _data_loader.create_torch_dataset_naive(data_config, action_horizon, model_config)   # * Using create_torch_dataset_naive

    dataset = _data_loader.TransformedDataset(
        dataset,
        [
            *data_config.repack_transforms.inputs,
            *data_config.data_transforms.inputs,
            # Remove strings since they are not supported by JAX and are not needed to compute norm stats.
            RemoveStrings(),
        ],
    )
    if max_frames is not None and max_frames < len(dataset):
        num_batches = max_frames // batch_size
        shuffle = True
    else:
        num_batches = len(dataset) // batch_size
        shuffle = False
    data_loader = _data_loader.TorchDataLoader(
        dataset,
        local_batch_size=batch_size,
        num_workers=num_workers,
        shuffle=shuffle,
        num_batches=num_batches,
    )
    return data_loader, num_batches


def main(config_name: str, max_frames: int | None = None):
    config = _config.get_config(config_name)
    data_config: _config.LerobotCustomAgilexDataConfig = config.data.create(config.assets_dirs, config.model)
    new_data_transforms = transforms.Group(
                inputs=[
                    FakeInputs(
                        action_dim=config.model.action_dim,
                        model_type=config.model.model_type,
                        mask_state=False,
                        convert_to_eef_position=False,
                    )
                ],
                outputs=[CustomAgilexOutputs()],
            )
    data_config = dataclasses.replace(
        data_config,
        repack_transforms=transforms.Group(
            inputs=[
                transforms.RepackTransform(
                    {
                        "state": "observation.state",
                        "actions": "action",
                    }
                )
            ]
        ),
        data_transforms=new_data_transforms,
    )

    assets_dir = pathlib.Path(config.data.assets.assets_dir) if config.data.assets.assets_dir else config.assets_dirs
    if data_config.asset_id is not None:
        output_path = assets_dir / data_config.asset_id
    else:
        repo_id = data_config.repo_id
        if isinstance(repo_id, list):
            if len(repo_id) != 1:
                raise ValueError("Need to specify assets.asset_id when using multiple datasets")
            repo_id = repo_id[0]
        output_path = assets_dir / pathlib.Path(repo_id).name

    print("Output_path:", output_path)

    
    data_loader, num_batches = create_torch_dataloader(
        data_config, config.model.action_horizon, config.batch_size, config.model, config.num_workers, max_frames
    )

    keys = ["state", "actions"]
    stats = {key: normalize.RunningStats() for key in keys}

    for batch in tqdm.tqdm(data_loader, total=num_batches, desc="Computing stats"):
        for key in keys:

            stats[key].update(np.asarray(batch[key]))

    norm_stats = {key: stats.get_statistics() for key, stats in stats.items()}

    print(f"Writing stats to: {output_path}")
    normalize.save(output_path, norm_stats)


if __name__ == "__main__":
    tyro.cli(main)