toddlerbot.manipulation.utils package¶
Submodules¶
toddlerbot.manipulation.utils.dataset_utils module¶
- toddlerbot.manipulation.utils.dataset_utils.create_sample_indices(episode_ends: ndarray, sequence_length: int, pad_before: int = 0, pad_after: int = 0)¶
Generates sample indices for sequences within episodes, considering padding.
- Parameters:
episode_ends (np.ndarray) – An array of indices indicating the end of each episode.
sequence_length (int) – The length of the sequence to sample.
pad_before (int, optional) – Number of padding steps to allow before the start of a sequence. Defaults to 0.
pad_after (int, optional) – Number of padding steps to allow after the end of a sequence. Defaults to 0.
- Returns:
An array of shape (n, 4) where each row contains the start and end indices of the buffer and the start and end indices of the sample within the sequence.
- Return type:
np.ndarray
- toddlerbot.manipulation.utils.dataset_utils.create_video_grid(image_data: ndarray, episode_ends: ndarray, save_path: str, file_name: str, num_cols: int = 10, fps: int = 10)¶
Creates a grid video from a sequence of image data, organizing episodes into a specified number of columns and saving the result to a file.
- Parameters:
image_data (np.ndarray) – A 4D numpy array of shape (T, C, H, W) containing the image data for all episodes, where T is the total number of frames, C is the number of channels, H is the height, and W is the width.
episode_ends (np.ndarray) – A 1D numpy array indicating the end frame indices for each episode.
save_path (str) – The directory path where the video file will be saved.
file_name (str) – The name of the video file to be created.
num_cols (int, optional) – The number of columns in the video grid. Defaults to 10.
fps (int, optional) – The frames per second for the output video. Defaults to 10.
- toddlerbot.manipulation.utils.dataset_utils.get_data_stats(data)¶
Calculates the minimum and maximum values for each feature in the given dataset.
- Parameters:
data (np.ndarray) – A multi-dimensional numpy array representing the dataset.
- Returns:
A dictionary containing the minimum and maximum values for each feature, with keys ‘min’ and ‘max’.
- Return type:
dict
- toddlerbot.manipulation.utils.dataset_utils.normalize_data(data, stats)¶
Normalizes the input data to a range of [-1, 1] based on provided statistics.
This function adjusts the input data by first normalizing it to a [0, 1] range using the minimum and maximum values from the stats dictionary. It then scales the normalized data to a [-1, 1] range. Indices where the range is zero are set to zero in the output.
- Parameters:
data (numpy.ndarray) – The input data to be normalized.
stats (dict) – A dictionary containing ‘min’ and ‘max’ keys with corresponding numpy arrays for minimum and maximum values.
- Returns:
The normalized data with values scaled to the range [-1, 1].
- Return type:
numpy.ndarray
- toddlerbot.manipulation.utils.dataset_utils.sample_sequence(train_data, sequence_length, buffer_start_idx, buffer_end_idx, sample_start_idx, sample_end_idx)¶
Generates a sequence sample from the given training data with specified indices and sequence length.
- Parameters:
train_data (dict) – A dictionary where keys are identifiers and values are numpy arrays representing the data.
sequence_length (int) – The desired length of the output sequence.
buffer_start_idx (int) – The starting index for slicing the input arrays.
buffer_end_idx (int) – The ending index for slicing the input arrays.
sample_start_idx (int) – The starting index for placing the sliced data in the output sequence.
sample_end_idx (int) – The ending index for placing the sliced data in the output sequence.
- Returns:
A dictionary with the same keys as train_data, where each value is a numpy array of the specified sequence length, containing the sampled data.
- Return type:
dict
- toddlerbot.manipulation.utils.dataset_utils.unnormalize_data(ndata, stats)¶
Converts normalized data back to its original scale using provided statistics.
- Parameters:
ndata (float or array-like) – The normalized data, assumed to be in the range [0, 1].
stats (dict) – A dictionary containing ‘max’ and ‘min’ keys with corresponding values representing the maximum and minimum of the original data.
- Returns:
The data rescaled to its original range.
- Return type:
float or array-like