Skip to content

Utils

Utility helpers for optimization workflows.

load_dataset

load_dataset(
    dataset_path,
    *,
    N=10,
    dt_hundredths=5,
    valid_states=(2, 4, 5),
    dist_threshold=100.0,
    vel_threshold=0.5,
    max_samples=None,
)

Load a single dataset folder into feature and simulation arrays.

load_feature_dataset_group

load_feature_dataset_group(
    dataset_paths,
    cache_path,
    refresh_cache,
    samples,
    dt_hundredths,
    label,
)

Load and optionally cache a list of dataset folders.

truncate_dataset_timesteps

truncate_dataset_timesteps(dataset, max_timesteps)

Trim controls/targets to the provided horizon (no-op if already shorter).

Dataset helpers used by optimization scripts.

load_dataset

load_dataset(
    dataset_path,
    *,
    N=10,
    dt_hundredths=5,
    valid_states=(2, 4, 5),
    dist_threshold=100.0,
    vel_threshold=0.5,
    max_samples=None,
)

Load a single dataset folder into feature and simulation arrays.

load_feature_dataset_group

load_feature_dataset_group(
    dataset_paths,
    cache_path,
    refresh_cache,
    samples,
    dt_hundredths,
    label,
)

Load and optionally cache a list of dataset folders.

truncate_dataset_timesteps

truncate_dataset_timesteps(dataset, max_timesteps)

Trim controls/targets to the provided horizon (no-op if already shorter).