| """ |
| Action normalization functions for robot datasets. |
| Provides unified [0, 1] range normalization. |
| """ |
|
|
| import torch |
| import numpy as np |
| import logging |
| from typing import Optional, Tuple |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def normalize_actions(actions: torch.Tensor, action_min: np.ndarray, action_max: np.ndarray) -> torch.Tensor: |
| """ |
| Normalize actions to [0, 1] range using provided statistics. |
| Formula: (action - min) / (max - min) |
| |
| Args: |
| actions: Action tensor to normalize |
| action_min: Minimum values for each action dimension |
| action_max: Maximum values for each action dimension |
| |
| Returns: |
| Normalized action tensor in [0, 1] range |
| """ |
| if action_min is None or action_max is None: |
| raise ValueError("Normalization stats (action_min/action_max) are None - cannot normalize actions") |
| |
| |
| actions_np = actions.numpy() if isinstance(actions, torch.Tensor) else actions |
| |
| |
| action_range = action_max - action_min |
| |
| action_range = np.where(action_range == 0, 1.0, action_range) |
| |
| normalized = (actions_np - action_min) / action_range |
| |
| return torch.from_numpy(normalized).float() |
|
|
|
|
| def denormalize_actions(normalized_actions: torch.Tensor, action_min: np.ndarray, action_max: np.ndarray) -> torch.Tensor: |
| """ |
| Denormalize actions from [0, 1] back to original scale. |
| Formula: normalized * (max - min) + min |
| |
| Args: |
| normalized_actions: Normalized action tensor in [0, 1] range |
| action_min: Minimum values for each action dimension |
| action_max: Maximum values for each action dimension |
| |
| Returns: |
| Denormalized action tensor in original scale |
| """ |
| if action_min is None or action_max is None: |
| raise ValueError("Normalization stats (action_min/action_max) are None - cannot denormalize actions") |
| |
| |
| norm_actions_np = normalized_actions.numpy() if isinstance(normalized_actions, torch.Tensor) else normalized_actions |
| |
| |
| action_range = action_max - action_min |
| denormalized = norm_actions_np * action_range + action_min |
| |
| return torch.from_numpy(denormalized).float() |
|
|
|
|
| def load_normalization_stats(stats_path: str, dataset_name: str) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: |
| """ |
| Load normalization statistics from file for specified dataset. |
| |
| Args: |
| stats_path: Path to normalization statistics file |
| dataset_name: Name of dataset to load stats for (e.g., 'ac_one', 'aloha_agilex_1') |
| |
| Returns: |
| Tuple of (action_min, action_max) arrays, or (None, None) if loading fails |
| """ |
| try: |
| import json |
| with open(stats_path, 'r') as f: |
| stats = json.load(f) |
| |
| |
| if dataset_name in stats: |
| dataset_stats = stats[dataset_name] |
| action_min = np.array(dataset_stats['min'], dtype=np.float32) |
| action_max = np.array(dataset_stats['max'], dtype=np.float32) |
| else: |
| raise KeyError(f"Dataset '{dataset_name}' not found in normalization stats file") |
| |
| logger.info(f"Loaded normalization stats for {dataset_name} from {stats_path}") |
| logger.info(f" Action min: {action_min}") |
| logger.info(f" Action max: {action_max}") |
| logger.info(f" Action range: {action_max - action_min}") |
| |
| return action_min, action_max |
| |
| except FileNotFoundError: |
| logger.warning(f"Normalization stats file not found: {stats_path}") |
| return None, None |
|
|
|
|
| def load_quantile_stats( |
| stats_path: str, |
| dataset_name: str, |
| lower_key: str = 'q01', |
| upper_key: str = 'q99', |
| ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: |
| """ |
| Load lower/upper quantile arrays for the specified dataset from a JSON file. |
| |
| Expected JSON structure example: |
| { |
| "latent_action": {"q01": [...], "q99": [...], ...} |
| } |
| |
| Args: |
| stats_path: Path to the quantile stats JSON file |
| dataset_name: Dataset name (e.g., 'latent_action') |
| lower_key: Lower quantile field name (default 'q01') |
| upper_key: Upper quantile field name (default 'q99') |
| |
| Returns: |
| Tuple (q_low, q_high) as np.ndarray; (None, None) on failure |
| """ |
| try: |
| import json |
| with open(stats_path, 'r') as f: |
| stats = json.load(f) |
|
|
| if dataset_name not in stats: |
| raise KeyError(f"Dataset '{dataset_name}' not found in stats file") |
| dataset_stats = stats[dataset_name] |
|
|
| if lower_key not in dataset_stats or upper_key not in dataset_stats: |
| raise KeyError( |
| f"Keys '{lower_key}'/'{upper_key}' not found in dataset '{dataset_name}' stats" |
| ) |
|
|
| q_low = np.array(dataset_stats[lower_key], dtype=np.float32) |
| q_high = np.array(dataset_stats[upper_key], dtype=np.float32) |
|
|
| logger.info( |
| f"Loaded quantile stats for {dataset_name} from {stats_path} ({lower_key}/{upper_key})" |
| ) |
| return q_low, q_high |
|
|
| except FileNotFoundError: |
| logger.warning(f"Quantile stats file not found: {stats_path}") |
| return None, None |
| except Exception as e: |
| logger.error(f"Error loading quantile stats from {stats_path}: {e}") |
| return None, None |
|
|
|
|
| def normalize_actions_with_quantiles( |
| actions: torch.Tensor, |
| q_low: np.ndarray, |
| q_high: np.ndarray, |
| *, |
| clip: bool = True, |
| ) -> torch.Tensor: |
| """ |
| Normalize to [0, 1] using quantiles: optional clipping to [q_low, q_high] then linear scaling. |
| |
| Formulas: |
| 1) if clip=True, \( x_{clip} = \min(\max(x, q_{01}), q_{99}) \) |
| 2) \( x_{norm} = (x_{clip} - q_{01}) / (q_{99} - q_{01}) \) |
| |
| Args: |
| actions: Action tensor [*, D] |
| q_low: Lower quantile array (e.g., q01), shape [D] |
| q_high: Upper quantile array (e.g., q99), shape [D] |
| clip: Whether to clip actions to [q_low, q_high] before scaling |
| |
| Returns: |
| Tensor normalized to [0, 1] |
| """ |
| if q_low is None or q_high is None: |
| raise ValueError("Quantile stats (q_low/q_high) are None - cannot normalize actions") |
|
|
| actions_np = actions.numpy() if isinstance(actions, torch.Tensor) else actions |
|
|
| |
| q_range = q_high - q_low |
| q_range = np.where(q_range == 0, 1.0, q_range) |
|
|
| x = actions_np |
| if clip: |
| x = np.clip(x, q_low, q_high) |
|
|
| normalized = (x - q_low) / q_range |
| return torch.from_numpy(normalized).float() |
|
|
|
|
