""" 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") # Convert to numpy for calculation actions_np = actions.numpy() if isinstance(actions, torch.Tensor) else actions # Normalize to [0, 1] range: (action - min) / (max - min) action_range = action_max - action_min # Avoid division by zero 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") # Convert to numpy for calculation norm_actions_np = normalized_actions.numpy() if isinstance(normalized_actions, torch.Tensor) else normalized_actions # Denormalize from [0, 1] back to original range: norm * (max - min) + min 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) # Get stats for specific dataset 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 # Avoid divide-by-zero 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()