"""Compute normalization statistics for BEHAVIOR-1K dataset. Computes mean, std, min, max from dataset in parallel. Supports per-timestamp normalization and action correlation matrices. """ import numpy as np import pandas as pd import tyro from pathlib import Path from concurrent.futures import ProcessPoolExecutor, as_completed import multiprocessing as mp import openpi.models.model as _model import openpi.transforms as transforms # Import B1K-specific modules from b1k.shared import normalize from b1k.training import config as _config from b1k.policies.b1k_policy import extract_state_from_proprio def get_delta_transform_from_config(config_name: str): """Get the delta action transform from the config.""" config = _config.get_config(config_name) data_config = config.data.create(config.assets_dirs, config.model) # Find the DeltaActions transform in the data transforms delta_transform = None for transform in data_config.data_transforms.inputs: if isinstance(transform, transforms.DeltaActions): delta_transform = transform break if delta_transform is None: raise ValueError("No DeltaActions transform found in config") return delta_transform def apply_delta_transform_from_config(state: np.ndarray, actions: np.ndarray, mask) -> np.ndarray: """Apply delta action transform using the mask from config. Args: state: Current state (23-dim) actions: Absolute actions (23-dim) mask: Boolean mask from DeltaActions transform Returns: delta_actions: Delta actions relative to current state (23-dim) """ if mask is None: return actions delta_actions = actions.copy() mask = np.asarray(mask) dims = mask.shape[-1] # Apply delta transform: delta = action - state for specified dimensions delta_actions[:dims] = np.where(mask, actions[:dims] - state[:dims], actions[:dims]) return delta_actions def process_episode_file(args): """Process a single episode file and return statistics.""" episode_file, delta_mask, action_horizon, compute_per_timestamp, compute_correlation, sample_fraction = args try: # Read parquet file directly df = pd.read_parquet(episode_file) # Extract states and actions states = [] raw_actions = [] # Keep raw actions for per-timestamp processing actions = [] for _, row in df.iterrows(): # Get raw proprioception and actions raw_state = np.array(row["observation.state"]) # 256-dim raw_action = np.array(row["action"]) # 23-dim # Apply state extraction (same as training/inference) processed_state = extract_state_from_proprio(raw_state) # 23-dim # Apply delta transform to actions (for regular statistics) delta_action = apply_delta_transform_from_config(processed_state, raw_action, delta_mask) states.append(processed_state) raw_actions.append(raw_action) actions.append(delta_action) if len(states) == 0: return None, None states = np.array(states) raw_actions = np.array(raw_actions) actions = np.array(actions) # Compute episode statistics (including min/max which we'll use instead of quantiles) episode_stats = { "state": { "count": len(states), "sum": np.sum(states, axis=0), "sum_sq": np.sum(states**2, axis=0), "min": np.min(states, axis=0), "max": np.max(states, axis=0), }, "actions": { "count": len(actions), "sum": np.sum(actions, axis=0), "sum_sq": np.sum(actions**2, axis=0), "min": np.min(actions, axis=0), "max": np.max(actions, axis=0), } } # Compute per-timestamp statistics and action chunks if requested per_timestamp_data = None correlation_chunks = None if compute_per_timestamp or compute_correlation: # Create action chunks for per-timestamp statistics # For each timestep, take the current state and next action_horizon absolute actions action_chunks = [] for i in range(len(states) - action_horizon + 1): current_state = states[i] # State at time i # Take next action_horizon absolute actions starting from time i future_absolute_actions = raw_actions[i:i + action_horizon] # Shape: (action_horizon, action_dim) # Apply delta transform to each action using the SAME current state delta_chunk = np.zeros_like(future_absolute_actions) for t in range(action_horizon): delta_chunk[t] = apply_delta_transform_from_config(current_state, future_absolute_actions[t], delta_mask) action_chunks.append(delta_chunk) if action_chunks: action_chunks = np.array(action_chunks) # Shape: (num_chunks, action_horizon, action_dim) # Sample chunks (not individual actions!) to reduce memory usage # Use fixed seed per episode for reproducibility rng = np.random.RandomState(hash(str(episode_file)) % (2**31)) n_chunk_samples = max(1, int(len(action_chunks) * sample_fraction)) if sample_fraction < 1.0 and n_chunk_samples < len(action_chunks): chunk_indices = rng.choice(len(action_chunks), size=n_chunk_samples, replace=False) sampled_chunks = action_chunks[chunk_indices] else: sampled_chunks = action_chunks if compute_per_timestamp: per_timestamp_data = sampled_chunks if compute_correlation: correlation_chunks = sampled_chunks # Don't return full states/actions to save memory - we have everything we need in episode_stats return episode_stats, (per_timestamp_data, correlation_chunks) except Exception as e: print(f"Error processing {episode_file}: {e}") return None, None def aggregate_episode_stats(episode_stats_list, all_data_list, config, compute_correlation=False, max_correlation_samples=2000000, compute_quantiles_sample_size=1000000): """Aggregate statistics from multiple episodes.""" # Collect per-timestamp data if available per_timestamp_chunks = [] for data in all_data_list: if data is not None and len(data) > 0 and data[0] is not None: per_timestamp_chunks.append(data[0]) # Collect correlation chunks if available correlation_chunks = [] for data in all_data_list: if data is not None and len(data) > 1 and data[1] is not None: correlation_chunks.append(data[1]) # Collect raw data for quantile computation (sample if too large) state_data_for_quantiles = [] action_data_for_quantiles = [] for data in all_data_list: if data is not None and len(data) > 0 and data[0] is not None: # data[0] is per_timestamp chunks: [N, H, D] chunks = data[0] # For actions, we have action chunks # For states, we would need separate collection (skip for now, use min/max) action_data_for_quantiles.append(chunks) # Aggregate counts, sums, etc. final_stats = {"state": {}, "actions": {}} for key in ["state", "actions"]: total_count = sum(stats[key]["count"] for stats in episode_stats_list if stats is not None) total_sum = np.sum([stats[key]["sum"] for stats in episode_stats_list if stats is not None], axis=0) total_sum_sq = np.sum([stats[key]["sum_sq"] for stats in episode_stats_list if stats is not None], axis=0) global_min = np.min([stats[key]["min"] for stats in episode_stats_list if stats is not None], axis=0) global_max = np.max([stats[key]["max"] for stats in episode_stats_list if stats is not None], axis=0) # Compute final statistics mean = total_sum / total_count variance = (total_sum_sq / total_count) - mean**2 std = np.sqrt(np.maximum(0, variance)) # Compute actual percentiles for actions (sample if needed for memory efficiency) if key == "actions" and action_data_for_quantiles: print(f"Computing actual percentiles for actions...") # Concatenate all action chunks and flatten to [N*H, D] all_action_data = np.vstack(action_data_for_quantiles) # [N, H, D] all_action_data_flat = all_action_data.reshape(-1, all_action_data.shape[-1]) # [N*H, D] # Sample if too large for memory if len(all_action_data_flat) > compute_quantiles_sample_size: print(f" Sampling {compute_quantiles_sample_size:,} from {len(all_action_data_flat):,} data points") rng = np.random.RandomState(42) indices = rng.choice(len(all_action_data_flat), size=compute_quantiles_sample_size, replace=False) all_action_data_flat = all_action_data_flat[indices] q01 = np.percentile(all_action_data_flat, 1, axis=0) q99 = np.percentile(all_action_data_flat, 99, axis=0) print(f" Computed q01/q99 from {len(all_action_data_flat):,} samples") else: # Fallback to min/max for state (no chunks collected) q01 = global_min q99 = global_max final_stats[key] = { "mean": mean, "std": std, "q01": q01, "q99": q99, } # Compute per-timestamp statistics for actions if available per_timestamp_stats = None if per_timestamp_chunks and key == "actions": # Combine all chunks: (num_episodes * num_chunks_per_episode, action_horizon, action_dim) all_chunks = np.vstack(per_timestamp_chunks) # Compute statistics for each timestamp action_horizon, action_dim = all_chunks.shape[1], all_chunks.shape[2] per_timestamp_mean = np.zeros((action_horizon, action_dim)) per_timestamp_std = np.zeros((action_horizon, action_dim)) per_timestamp_q01 = np.zeros((action_horizon, action_dim)) per_timestamp_q99 = np.zeros((action_horizon, action_dim)) for t in range(action_horizon): timestep_data = all_chunks[:, t, :] # Shape: (num_chunks, action_dim) per_timestamp_mean[t] = np.mean(timestep_data, axis=0) per_timestamp_std[t] = np.std(timestep_data, axis=0) per_timestamp_q01[t] = np.percentile(timestep_data, 1, axis=0) per_timestamp_q99[t] = np.percentile(timestep_data, 99, axis=0) per_timestamp_stats = { "per_timestamp_mean": per_timestamp_mean, "per_timestamp_std": per_timestamp_std, "per_timestamp_q01": per_timestamp_q01, "per_timestamp_q99": per_timestamp_q99, } # Compute full correlation matrix for actions if available correlation_stats = None if compute_correlation and not correlation_chunks: raise RuntimeError( "Correlation computation was requested but no correlation chunks were collected. " "This indicates an issue with episode processing or insufficient data." ) if correlation_chunks: print("Computing full action correlation matrix...") # Combine all chunks: (num_episodes * num_chunks_per_episode, action_horizon, action_dim) all_corr_chunks = np.vstack(correlation_chunks) action_horizon, action_dim = all_corr_chunks.shape[1], all_corr_chunks.shape[2] target_action_dim = config.model.action_dim # 32 (with padding) print(f"Total action chunks available: {all_corr_chunks.shape[0]} of shape ({action_horizon}, {action_dim})") # PAD action dimensions BEFORE computing correlations if action_dim < target_action_dim: print(f"Padding actions from {action_dim}D to {target_action_dim}D...") # Pad with zeros: (num_samples, action_horizon, action_dim) -> (num_samples, action_horizon, target_action_dim) padding = np.zeros((all_corr_chunks.shape[0], action_horizon, target_action_dim - action_dim)) all_corr_chunks_padded = np.concatenate([all_corr_chunks, padding], axis=2) print(f"Padded chunks shape: {all_corr_chunks_padded.shape}") else: all_corr_chunks_padded = all_corr_chunks[:, :, :target_action_dim] # Flatten chunks to (num_samples, action_horizon * target_action_dim) flattened_chunks = all_corr_chunks_padded.reshape(-1, action_horizon * target_action_dim) total_samples = flattened_chunks.shape[0] # Subsample if we have too many samples (for computational efficiency) if total_samples > max_correlation_samples: print(f"Subsampling {max_correlation_samples} random chunks for correlation computation (from {total_samples})") print(f"This provides ~{100 * max_correlation_samples / total_samples:.1f}x speedup with negligible accuracy loss") rng = np.random.RandomState(42) # Fixed seed for reproducibility sample_indices = rng.choice(total_samples, size=max_correlation_samples, replace=False) flattened_chunks = flattened_chunks[sample_indices] print(f"Using {flattened_chunks.shape[0]} sampled chunks for correlation matrix computation") else: print(f"Using all {total_samples} chunks (no subsampling needed)") # Normalize each dimension to zero mean and unit variance for correlation computation # This ensures the correlation matrix captures pure correlation structure chunk_mean = np.mean(flattened_chunks, axis=0) chunk_std = np.std(flattened_chunks, axis=0) # Identify constant dimensions (std ≈ 0, including padded dimensions) constant_dims = chunk_std < 1e-6 print(f"Found {np.sum(constant_dims)} constant/padded dimensions (will be set to identity)") # Normalize non-constant dimensions normalized_chunks = flattened_chunks.copy() normalized_chunks[:, ~constant_dims] = (flattened_chunks[:, ~constant_dims] - chunk_mean[~constant_dims]) / chunk_std[~constant_dims] # Constant dimensions stay as 0 (already centered at 0 due to padding) # Compute empirical covariance matrix # For normalized data, covariance = correlation cov_matrix = np.cov(normalized_chunks, rowvar=False) # Shape: [H*D, H*D] print(f"Covariance matrix shape: {cov_matrix.shape}") # Enforce diagonal = 1 for all dimensions (standard correlation matrix property) # Handle constant dimensions carefully (avoid division by zero) diag_vals = np.diag(cov_matrix).copy() print(f"Diagonal before correction: min={np.min(diag_vals):.6f}, max={np.max(diag_vals):.6f}") # Identify truly constant dimensions (diagonal ≈ 0) constant_mask = diag_vals < 1e-10 num_constant = np.sum(constant_mask) print(f"Found {num_constant} dimensions with zero variance on diagonal") # Normalize correlation matrix (avoid division by zero) diag_vals_safe = diag_vals.copy() diag_vals_safe[constant_mask] = 1.0 # Prevent division by zero normalizer = np.sqrt(diag_vals_safe[:, None] @ diag_vals_safe[None, :]) cov_matrix = cov_matrix / normalizer # For constant dimensions: set their rows/columns to 0, diagonal to 1 cov_matrix[constant_mask, :] = 0.0 cov_matrix[:, constant_mask] = 0.0 np.fill_diagonal(cov_matrix, 1.0) print(f"Diagonal after correction: min={np.min(np.diag(cov_matrix)):.6f}, max={np.max(np.diag(cov_matrix)):.6f}") print(f"Matrix check: has NaN={np.any(np.isnan(cov_matrix))}, has Inf={np.any(np.isinf(cov_matrix))}") # Add regularization for numerical stability epsilon = 1e-6 cov_matrix_reg = cov_matrix + epsilon * np.eye(cov_matrix.shape[0]) # Check eigenvalues for positive definiteness eigenvalues = np.linalg.eigvalsh(cov_matrix_reg) min_eigenvalue = np.min(eigenvalues) print(f"Min eigenvalue after regularization: {min_eigenvalue:.6e}") if min_eigenvalue <= 0: print(f"WARNING: Matrix not positive definite! Adding stronger regularization...") # Add stronger regularization if needed epsilon = max(1e-5, -min_eigenvalue + 1e-5) cov_matrix_reg = cov_matrix + epsilon * np.eye(cov_matrix.shape[0]) eigenvalues = np.linalg.eigvalsh(cov_matrix_reg) min_eigenvalue = np.min(eigenvalues) print(f"New min eigenvalue: {min_eigenvalue:.6e}") # Compute Cholesky decomposition try: chol_lower = np.linalg.cholesky(cov_matrix_reg) print(f"Cholesky decomposition successful! Shape: {chol_lower.shape}") # Verify reconstruction: L @ L.T ≈ Σ reconstructed = chol_lower @ chol_lower.T reconstruction_error = np.linalg.norm(reconstructed - cov_matrix_reg, 'fro') / np.linalg.norm(cov_matrix_reg, 'fro') print(f"Cholesky reconstruction error: {reconstruction_error:.6e}") # Compute averaged spatial correlation (dim × dim) # Average correlation across all timesteps print("\nComputing averaged spatial correlation (dim × dim)...") spatial_corrs = [] for t in range(action_horizon): start_idx = t * target_action_dim end_idx = (t + 1) * target_action_dim timestep_data = normalized_chunks[:, start_idx:end_idx] corr_t = np.cov(timestep_data, rowvar=False) # Normalize to ensure diagonal = 1 diag_t = np.diag(corr_t) corr_t = corr_t / np.sqrt(diag_t[:, None] @ diag_t[None, :]) spatial_corrs.append(corr_t) avg_spatial_corr = np.mean(spatial_corrs, axis=0) # (target_action_dim, target_action_dim) # Final diagonal enforcement np.fill_diagonal(avg_spatial_corr, 1.0) print(f"Averaged spatial correlation shape: {avg_spatial_corr.shape}") # Compute averaged temporal correlation (time × time) # Average correlation across all NON-CONSTANT dimensions only print("Computing averaged temporal correlation (time × time)...") temporal_corrs = [] # Identify which dimensions are constant (per original dimension, before flattening) dim_is_constant = [] for d in range(target_action_dim): # Check if this dimension is constant across all timesteps dim_indices_check = [t * target_action_dim + d for t in range(action_horizon)] dim_std = chunk_std[dim_indices_check] # Dimension is constant if std < 1e-6 for all its timesteps dim_is_constant.append(np.all(dim_std < 1e-6)) num_constant = np.sum(dim_is_constant) print(f" Excluding {num_constant} constant dimensions from temporal averaging") for d in range(target_action_dim): if dim_is_constant[d]: continue # Skip constant dimensions dim_indices = [t * target_action_dim + d for t in range(action_horizon)] dim_data = normalized_chunks[:, dim_indices] corr_d = np.cov(dim_data, rowvar=False) # Normalize to ensure diagonal = 1 diag_d = np.diag(corr_d) corr_d = corr_d / np.sqrt(diag_d[:, None] @ diag_d[None, :]) temporal_corrs.append(corr_d) if len(temporal_corrs) > 0: avg_temporal_corr = np.mean(temporal_corrs, axis=0) # (action_horizon, action_horizon) # Final diagonal enforcement np.fill_diagonal(avg_temporal_corr, 1.0) print(f"Averaged temporal correlation shape: {avg_temporal_corr.shape}") else: # All dimensions are constant - use identity avg_temporal_corr = np.eye(action_horizon) print(f"All dimensions constant - using identity for temporal correlation") # Store correlation matrices # Beta shrinkage will be applied at model load time correlation_stats = { "action_correlation_cholesky": chol_lower, "action_correlation_spatial": avg_spatial_corr, # (D, D) - averaged spatial "action_correlation_temporal": avg_temporal_corr, # (H, H) - averaged temporal } except np.linalg.LinAlgError as e: raise RuntimeError( f"Cholesky decomposition failed: {e}. " "This indicates the covariance matrix is not positive definite even after regularization. " "This is a critical error that needs investigation." ) return final_stats, per_timestamp_stats, correlation_stats def main( config_name: str, max_episodes: int | None = None, num_workers: int | None = None, per_timestamp: bool = False, correlation: bool = True, max_correlation_samples: int = 2000000, sample_fraction: float = 0.1 ): """Compute normalization statistics quickly for B1K config using parallel processing. Args: config_name: Name of the training config max_episodes: Maximum number of episodes to process (None = all) num_workers: Number of parallel workers (None = auto) per_timestamp: Whether to compute per-timestamp statistics correlation: Whether to compute action correlation matrix max_correlation_samples: Maximum number of samples to use for correlation computation (subsampling for efficiency, 2M is typically sufficient) sample_fraction: Fraction of action chunks to sample from each episode for correlation/per-timestamp stats (default 0.1 = 10%, use 1.0 for no sampling). Note: mean/std/min/max are ALWAYS computed over ALL data. """ config = _config.get_config(config_name) data_config = config.data.create(config.assets_dirs, config.model) if not data_config.behavior_dataset_root: raise ValueError("This script only works with B1K behavior datasets") # Get the delta transform from config delta_transform = get_delta_transform_from_config(config_name) print(f"Using delta transform with mask: {delta_transform.mask}") # Check if we should compute per-timestamp stats compute_per_timestamp = per_timestamp or data_config.use_per_timestamp_norm compute_correlation = correlation action_horizon = config.model.action_horizon if compute_per_timestamp: print(f"Computing per-timestamp statistics with action_horizon={action_horizon}") if compute_correlation: print(f"Computing correlation matrix for action chunks with action_horizon={action_horizon}") # Print sampling information print(f"Memory-efficient mode: Using min/max instead of quantiles (q01/q99)") if sample_fraction < 1.0: print(f"Sampling {sample_fraction*100:.0f}% of action chunks from each episode for correlation/per-timestamp stats") else: print(f"Using all action chunks (no sampling) - this may cause OOM on large datasets") # Find all episode parquet files data_root = Path(data_config.behavior_dataset_root) print(f"Looking for episode files in: {data_root}/data/task-*/episode_*.parquet") episode_files = list(data_root.glob("data/task-*/episode_*.parquet")) print(f"Found {len(episode_files)} total episode files") if max_episodes is not None: episode_files = episode_files[:max_episodes] print(f"Limited to {len(episode_files)} files based on max_episodes") if len(episode_files) == 0: print("No episode files found! Checking directory structure...") print(f"Contents of {data_root}:") for item in data_root.iterdir(): print(f" {item}") if (data_root / "data").exists(): print(f"Contents of {data_root}/data:") for item in (data_root / "data").iterdir(): print(f" {item}") raise ValueError("No episode files found") print(f"Processing {len(episode_files)} episode files...") # Set up parallel processing if num_workers is None: num_workers = min(mp.cpu_count(), len(episode_files)) # Ensure at least 1 worker num_workers = max(1, num_workers) print(f"Using {num_workers} workers") # Process episodes in parallel episode_stats_list = [] all_data_list = [] with ProcessPoolExecutor(max_workers=num_workers) as executor: # Submit all jobs future_to_file = { executor.submit(process_episode_file, (episode_file, delta_transform.mask, action_horizon, compute_per_timestamp, compute_correlation, sample_fraction)): episode_file for episode_file in episode_files } # Collect results for future in as_completed(future_to_file): episode_file = future_to_file[future] try: episode_stats, episode_data = future.result() if episode_stats is not None: episode_stats_list.append(episode_stats) all_data_list.append(episode_data) if len(episode_stats_list) % 10 == 0: print(f"Processed {len(episode_stats_list)} episodes...") except Exception as e: print(f"Error processing {episode_file}: {e}") print(f"Successfully processed {len(episode_stats_list)} episodes") # Aggregate all statistics print("Aggregating statistics...") final_stats, per_timestamp_stats, correlation_stats = aggregate_episode_stats( episode_stats_list, all_data_list, config, compute_correlation=compute_correlation, max_correlation_samples=max_correlation_samples ) # Prepare correlation matrix if computed correlation_matrix = None if compute_correlation: if correlation_stats is None: raise RuntimeError( "Correlation computation was requested (--correlation) but correlation_stats is None. " "This indicates an unexpected error in the correlation computation pipeline." ) print("Adding action correlation matrix...") chol_matrix = correlation_stats["action_correlation_cholesky"] print(f"Correlation matrix shape: {chol_matrix.shape}") # Verify size matches expected (should already be correct from padding before computation) expected_dim = action_horizon * config.model.action_dim if chol_matrix.shape[0] != expected_dim: raise ValueError( f"Correlation matrix has unexpected size {chol_matrix.shape[0]}, " f"expected {expected_dim} (action_horizon={action_horizon} × action_dim={config.model.action_dim})" ) correlation_matrix = chol_matrix elif correlation_stats is not None: # This shouldn't happen but handle it gracefully print("WARNING: Correlation matrix was computed but --correlation flag was False. Ignoring correlation matrix.") # Create NormStats objects with proper structure norm_stats_dict = {} for key in ["state", "actions"]: # Prepare basic statistics stats_kwargs = {} for stat_type in ["mean", "std", "q01", "q99"]: stat_value = final_stats[key][stat_type] padded_value = transforms.pad_to_dim(stat_value, config.model.action_dim) stats_kwargs[stat_type] = padded_value # Add per-timestamp statistics if computed and this is actions if per_timestamp_stats is not None and key == "actions": print("Adding per-timestamp statistics...") for stat_type in ["per_timestamp_mean", "per_timestamp_std", "per_timestamp_q01", "per_timestamp_q99"]: stat_value = per_timestamp_stats[stat_type] # Pad each timestamp to model action dimension padded_value = [] for t in range(stat_value.shape[0]): padded_timestep = transforms.pad_to_dim(stat_value[t], config.model.action_dim) padded_value.append(padded_timestep) stats_kwargs[stat_type] = np.array(padded_value) # Add correlation matrices if computed and this is actions if correlation_matrix is not None and key == "actions": stats_kwargs["action_correlation_cholesky"] = correlation_matrix # Also add averaged spatial and temporal correlations if available if correlation_stats and "action_correlation_spatial" in correlation_stats: stats_kwargs["action_correlation_spatial"] = correlation_stats["action_correlation_spatial"] if correlation_stats and "action_correlation_temporal" in correlation_stats: stats_kwargs["action_correlation_temporal"] = correlation_stats["action_correlation_temporal"] # Create NormStats object norm_stats_dict[key] = normalize.NormStats(**stats_kwargs) # Save statistics output_path = config.assets_dirs / data_config.repo_id output_path.mkdir(parents=True, exist_ok=True) normalize.save(output_path, norm_stats_dict) # Print summary for verification print("\nStatistics Summary (first 23 dims):") state_means = np.array(norm_stats_dict["state"].mean[:23]) action_means = np.array(norm_stats_dict["actions"].mean[:23]) differences = np.abs(state_means - action_means) print("State means: ", state_means) print("Action means: ", action_means) print("Differences: ", differences) print("Max difference:", np.max(differences)) print("\nNote: For actions, q01/q99 are actual percentiles. For state, q01/q99 contain min/max values.") if per_timestamp_stats is not None: print(f"\nPer-timestamp statistics computed for {action_horizon} timesteps") # Show first few timesteps for verification per_ts_means = np.array(norm_stats_dict["actions"].per_timestamp_mean) print(f"Per-timestamp means shape: {per_ts_means.shape}") print("First 3 timesteps, first 5 dims:") for t in range(min(3, per_ts_means.shape[0])): print(f" t={t}: {per_ts_means[t, :5]}") if compute_correlation: print(f"\nCorrelation matrices computed!") # Full correlation matrix corr_matrix = np.array(norm_stats_dict["actions"].action_correlation_cholesky) print(f" Full correlation (Cholesky): {corr_matrix.shape}") print(f" Memory: {corr_matrix.nbytes / 1024 / 1024:.2f} MB") # Averaged spatial correlation if hasattr(norm_stats_dict["actions"], 'action_correlation_spatial') and \ norm_stats_dict["actions"].action_correlation_spatial is not None: spatial = np.array(norm_stats_dict["actions"].action_correlation_spatial) print(f" Spatial (dim×dim, averaged): {spatial.shape}") print(f" Memory: {spatial.nbytes / 1024:.2f} KB") print(f" Sample correlation (d0 vs d1): {spatial[0, 1]:.4f}") # Averaged temporal correlation if hasattr(norm_stats_dict["actions"], 'action_correlation_temporal') and \ norm_stats_dict["actions"].action_correlation_temporal is not None: temporal = np.array(norm_stats_dict["actions"].action_correlation_temporal) print(f" Temporal (time×time, averaged): {temporal.shape}") print(f" Memory: {temporal.nbytes / 1024:.2f} KB") print(f" Sample correlation (t0 vs t1): {temporal[0, 1]:.4f}") print("\nNote: Beta shrinkage will be applied at model load time (not during stats computation)") else: print("\nCorrelation matrix computation was disabled (--no-correlation)") print(f"\nStatistics saved to: {output_path}") if __name__ == "__main__": tyro.cli(main)