opentau.datasets.compute_stats

Statistics computation and aggregation for dataset features.

This module provides functionality to compute statistical measures (min, max, mean, standard deviation, and count) for dataset features, with special handling for image and video data. It supports per-episode statistics computation and aggregation across multiple episodes or datasets using weighted averaging.

The module handles two main use cases:

Computing statistics for individual episodes: Samples images efficiently, downsamples large images to reduce memory usage, and computes statistics for all feature types (images, vectors, etc.).
Aggregating statistics across multiple episodes/datasets: Combines statistics using weighted mean and variance computation, taking global min/max values.

Key Features:

Memory-efficient image sampling: Uses heuristic-based sampling to estimate optimal number of samples based on dataset size.
Automatic image downsampling: Reduces large images (>300px) to ~150px for faster processing.
Weighted aggregation: Supports custom weights or uses episode counts as weights for aggregating statistics.
Parallel variance algorithm: Uses efficient algorithm for computing weighted variance across multiple statistics.

Functions:

estimate_num_samples: Heuristic to estimate optimal number of samples based on dataset size.
sample_indices: Generate evenly spaced sample indices from a dataset.
auto_downsample_height_width: Automatically downsample large images.
sample_images: Load and downsample a subset of images from file paths.
get_feature_stats: Compute statistical measures for an array.
compute_episode_stats: Compute statistics for a single episode.
aggregate_feature_stats: Aggregate statistics for a feature across multiple episodes.
aggregate_stats: Aggregate statistics from multiple episodes/datasets.

Example

Compute statistics for a single episode:

>>> episode_data = {"state": state_array, "camera0": image_paths}
>>> features = {"state": {"dtype": "float32"}, "camera0": {"dtype": "image"}}
>>> stats = compute_episode_stats(episode_data, features)

Aggregate statistics across multiple episodes:

>>> stats_list = [episode1_stats, episode2_stats, episode3_stats]
>>> weights = [100, 200, 150]  # Optional: custom weights
>>> aggregated = aggregate_stats(stats_list, weights=weights)

Functions

`aggregate_feature_stats`(stats_ft_list[, weights])	Aggregate statistics for a single feature across multiple episodes/datasets.
`aggregate_stats`(stats_list[, weights])	Aggregate stats from multiple compute_stats outputs into a single set of stats.
`auto_downsample_height_width`(img[, ...])	Automatically downsample an image if it exceeds size threshold.
`compute_episode_stats`(episode_data, features)	Compute statistics for a single episode.
`estimate_num_samples`(dataset_len[, ...])	Heuristic to estimate the number of samples based on dataset size.
`get_feature_stats`(array, axis, keepdims)	Compute statistical measures (min, max, mean, std, count) for an array.
`sample_images`(image_paths)	Load and downsample a subset of images from file paths.
`sample_indices`(data_len)	Generate evenly spaced sample indices from a dataset.

opentau.datasets.compute_stats.aggregate_feature_stats(stats_ft_list: list[dict[str, dict]], weights: list[float] | None = None) → dict[str, dict[str, ndarray]][source]

Aggregate statistics for a single feature across multiple episodes/datasets.

Computes weighted mean and variance using the parallel algorithm for variance computation. Min and max are taken as the global min/max across all stats.

Non-finite-tolerant per-dim: a contributor whose mean/std/min/max is non-finite (NaN or +/-Inf) at some dim is excluded from the aggregation at that dim only. Clean dims aggregate normally. A dim that is non-finite across every contributor stays NaN, so the downstream loader (or Normalize buffer) can surface it – but a single bad dataset no longer poisons the global buffer for every other sample.

The mean and variance use different per-dim masks: a contributor is dropped from the mean at a dim iff its mean is non-finite there, but dropped from the variance at a dim iff either its mean or its std is non-finite there (you can’t form a Chan-style variance contribution without a finite std). count is unconditional: the sum of contributors’ counts, not a per-dim effective contributor count.

Parameters:

stats_ft_list – List of statistics dictionaries for the same feature.
weights – Optional weights for each statistics entry. If None, uses count values as weights.

Returns:

Aggregated statistics dictionary with min, max, mean, std, and count.

opentau.datasets.compute_stats.aggregate_stats(stats_list: list[dict[str, dict]], weights: list[float] | None = None) → dict[str, dict[str, ndarray]][source]

Aggregate stats from multiple compute_stats outputs into a single set of stats.

The final stats will have the union of all data keys from each of the stats dicts.

For instance: - new_min = min(min_dataset_0, min_dataset_1, …) - new_max = max(max_dataset_0, max_dataset_1, …) - new_mean = (mean of all data, weighted by counts) - new_std = (std of all data)

opentau.datasets.compute_stats.auto_downsample_height_width(img: ndarray, target_size: int = 150, max_size_threshold: int = 300) → ndarray[source]

Automatically downsample an image if it exceeds size threshold.

If the image’s maximum dimension is below the threshold, returns it unchanged. Otherwise, downsamples by an integer factor to bring the larger dimension close to the target size.

Parameters:

img – Input image array of shape (C, H, W).
target_size – Target size for the larger dimension after downsampling. Defaults to 150.
max_size_threshold – Maximum size before downsampling is applied. Defaults to 300.

Returns:

Downsampled image array, or original if no downsampling needed.

opentau.datasets.compute_stats.compute_episode_stats(episode_data: dict[str, list[str] | ndarray], features: dict, skip_video_stats: bool = False) → dict[source]

Compute statistics for a single episode.

For image/video features, samples and downsamples images before computing stats (unless skip_video_stats is True, in which case placeholder stats are used). For other features, computes stats directly on the array data.

Parameters:

episode_data – Dictionary mapping feature names to their data (arrays or image paths).
features – Dictionary of feature specifications with ‘dtype’ keys.
skip_video_stats – If True, do not compute real stats for image/video features; instead use placeholder stats (min=0, max=1, mean=0.5, std=0.5, count from data) so the output format remains valid.

Returns:

Dictionary mapping feature names to their statistics (min, max, mean, std, count). Image statistics are normalized to [0, 1] range (or placeholders when skip_video_stats).

opentau.datasets.compute_stats.estimate_num_samples(dataset_len: int, min_num_samples: int = 100, max_num_samples: int = 10000, power: float = 0.75) → int[source]

Heuristic to estimate the number of samples based on dataset size. The power controls the sample growth relative to dataset size. Lower the power for less number of samples.

For default arguments, we have: - from 1 to ~500, num_samples=100 - at 1000, num_samples=177 - at 2000, num_samples=299 - at 5000, num_samples=594 - at 10000, num_samples=1000 - at 20000, num_samples=1681

opentau.datasets.compute_stats.get_feature_stats(array: ndarray, axis: tuple, keepdims: bool) → dict[str, ndarray][source]

Compute statistical measures (min, max, mean, std, count) for an array.

Parameters:

array – Input numpy array to compute statistics over.
axis – Axes along which to compute statistics.
keepdims – Whether to keep reduced dimensions.

Returns:

Dictionary containing ‘min’, ‘max’, ‘mean’, ‘std’, and ‘count’ statistics.

opentau.datasets.compute_stats.sample_images(image_paths: list[str]) → ndarray[source]

Load and downsample a subset of images from file paths.

Samples images using evenly spaced indices, loads them as uint8 arrays, and automatically downsamples large images to reduce memory usage.

Parameters:: image_paths – List of file paths to image files.
Returns:: Array of shape (num_samples, C, H, W) containing sampled images as uint8.

opentau.datasets.compute_stats.sample_indices(data_len: int) → list[int][source]

Generate evenly spaced sample indices from a dataset.

Uses estimate_num_samples to determine how many samples to take, then returns evenly spaced indices across the dataset length.

Parameters:: data_len – Total length of the dataset.
Returns:: List of evenly spaced integer indices.