| import abc |
| import utils |
| import torch |
| import numpy as np |
| from torch import default_generator, randperm |
| from torch.utils.data import Dataset, Subset |
| from typing import Callable, Optional, Sequence, List |
| from torch.nn.utils.rnn import pad_sequence |
|
|
|
|
| |
| def _accumulate(iterable, fn=lambda x, y: x + y): |
| "Return running totals" |
| |
| |
| it = iter(iterable) |
| try: |
| total = next(it) |
| except StopIteration: |
| return |
| yield total |
| for element in it: |
| total = fn(total, element) |
| yield total |
|
|
|
|
| class TrajectoryDataset(Dataset, abc.ABC): |
| """ |
| A dataset containing trajectories. |
| TrajectoryDataset[i] returns: (observations, actions, mask) |
| observations: Tensor[T, ...], T frames of observations |
| actions: Tensor[T, ...], T frames of actions |
| mask: Tensor[T]: False: invalid; True: valid |
| """ |
|
|
| @abc.abstractmethod |
| def get_seq_length(self, idx): |
| """ |
| Returns the length of the idx-th trajectory. |
| """ |
| raise NotImplementedError |
|
|
| @abc.abstractmethod |
| def get_frames(self, idx, frames): |
| """ |
| Returns the frames from the idx-th trajectory at the specified frames. |
| Used to speed up slicing. |
| """ |
| raise NotImplementedError |
|
|
|
|
| class TrajectorySubset(TrajectoryDataset, Subset): |
| """ |
| Subset of a trajectory dataset at specified indices. |
| |
| Args: |
| dataset (TrajectoryDataset): The whole Dataset |
| indices (sequence): Indices in the whole set selected for subset |
| """ |
|
|
| def __init__(self, dataset: TrajectoryDataset, indices: Sequence[int]): |
| Subset.__init__(self, dataset, indices) |
|
|
| def get_seq_length(self, idx): |
| return self.dataset.get_seq_length(self.indices[idx]) |
|
|
| def get_all_actions(self): |
| return self.dataset.get_all_actions() |
|
|
| def get_frames(self, idx, frames): |
| return self.dataset.get_frames(self.indices[idx], frames) |
|
|
|
|
| class TrajectorySlicerDataset(Dataset): |
| """ |
| Slice a trajectory dataset into (overlapping) windows of `window` observations |
| paired with an action chunk. |
| |
| dataset: a trajectory dataset that satisfies: |
| dataset.get_seq_length(i) returns the length of sequence i |
| dataset[i] = (observations, actions, *others) |
| observations: Tensor[T, ...] |
| actions: Tensor[T, ...] |
| window: int |
| number of observation timesteps in each slice |
| action_window: int |
| number of action timesteps to predict |
| vqbet_get_future_action_chunk: bool = True |
| if True, return only the action chunk following the observation window; |
| otherwise return actions spanning the whole window plus the chunk |
| transform: function (values) -> values |
| pad_seq_length: bool = True |
| pad actions at the end to ensure a fixed length instead of dropping short slices |
| |
| Goal conditioning is not handled here: any goal tensor is carried through as part |
| of `*others` by the underlying dataset. The remaining arguments |
| (future_conditional, min_future_sep, future_seq_len, only_sample_tail, |
| use_libero_goal) are recorded on the instance but do not affect slicing; |
| future_conditional only changes the value reported by get_seq_length. |
| """ |
|
|
| def __init__( |
| self, |
| dataset: TrajectoryDataset, |
| window: int, |
| action_window: int, |
| vqbet_get_future_action_chunk: bool = True, |
| future_conditional: bool = False, |
| min_future_sep: int = 0, |
| future_seq_len: Optional[int] = None, |
| only_sample_tail: bool = False, |
| transform: Optional[Callable] = None, |
| use_libero_goal: bool = False, |
| pad_seq_length: bool = True, |
| ): |
| if future_conditional: |
| assert future_seq_len is not None, "must specify a future_seq_len" |
| self.dataset = dataset |
| self.window = window |
| self.action_window = action_window |
| self.vqbet_get_future_action_chunk = vqbet_get_future_action_chunk |
| self.future_conditional = future_conditional |
| self.min_future_sep = min_future_sep |
| self.future_seq_len = future_seq_len |
| self.only_sample_tail = only_sample_tail |
| self.transform = transform |
| self.slices = [] |
| self.use_libero_goal = use_libero_goal |
| self.pad_seq_length = pad_seq_length |
|
|
| min_seq_length = np.inf |
| min_window_required = window + action_window - 1 |
| for i in range(len(self.dataset)): |
| T = self.dataset.get_seq_length(i) |
| min_seq_length = min(T, min_seq_length) |
|
|
| self.slices += [ |
| (i, 0, end + 1) for end in range(window - 1) |
| ] |
|
|
| if self.pad_seq_length: |
| if T - self.window >= 0: |
| self.slices += [ |
| (i, start, start + self.window) |
| for start in range(T - self.window + 1) |
| ] |
| else: |
| if T - min_window_required < 0: |
| print( |
| f"Ignored short sequence #{i}: len={T}, window={min_window_required}" |
| ) |
| else: |
| self.slices += [ |
| (i, start, start + self.window) |
| for start in range(T - min_window_required + 1) |
| ] |
|
|
| if (not self.pad_seq_length) and (min_seq_length < min_window_required): |
| print( |
| f"Ignored short sequences. To include all, set window <= {min_seq_length}." |
| ) |
|
|
| def get_seq_length(self, idx: int) -> int: |
| if self.future_conditional: |
| return self.future_seq_len + self.window |
| else: |
| return self.window |
|
|
| def __len__(self): |
| return len(self.slices) |
|
|
| def __getitem__(self, idx): |
| i, start, end = self.slices[idx] |
| obs, act, *others = self.dataset[i] |
| if end - start < self.window: |
| obs_win = utils.inference.repeat_start_to_length( |
| obs[start:end], self.window, dim=0 |
| ) |
| act = utils.inference.repeat_start_to_length( |
| |
| act[start : end - 1 + self.action_window], |
| self.window + self.action_window - 1, |
| dim=0, |
| ) |
| repeated_others = [ |
| utils.inference.repeat_start_to_length( |
| other[start:end], self.window, dim=0 |
| ) |
| for other in others |
| ] |
| else: |
| obs_win = obs[start:end] |
| |
| act = act[start : end - 1 + self.action_window] |
| repeated_others = [other[start:end] for other in others] |
|
|
| if self.vqbet_get_future_action_chunk: |
| expected_len = self.action_window |
| act = act[self.window - 1 :] |
| else: |
| expected_len = self.window + self.action_window - 1 |
|
|
| if act.shape[0] < expected_len: |
| if self.pad_seq_length: |
| act = utils.inference.repeat_end_to_length(act, expected_len, dim=0) |
| else: |
| raise ValueError( |
| f"Action chunk too short: {act.shape[0]} < {expected_len}, " |
| f"but pad_seq_length is False" |
| ) |
| values = [obs_win, act, *repeated_others] |
|
|
| |
| if self.transform is not None: |
| values = self.transform(values) |
| return tuple(values) |
|
|
|
|
| class TrajectoryEmbeddingDataset(TrajectoryDataset): |
| def __init__( |
| self, |
| model, |
| dataset: TrajectoryDataset, |
| device="cpu", |
| embed_goal=False, |
| ): |
| self.data = utils.inference.embed_trajectory_dataset( |
| model, |
| dataset, |
| obs_only=False, |
| device=device, |
| embed_goal=embed_goal, |
| ) |
| assert len(self.data) == len(dataset) |
|
|
| self.seq_lengths = [len(x[0]) for x in self.data] |
| self.on_device_data = [] |
| n_tensors = len(self.data[0]) |
| for i in range(n_tensors): |
| self.on_device_data.append( |
| pad_sequence([x[i] for x in self.data], batch_first=True).to(device) |
| ) |
| self.data = self.on_device_data |
|
|
| def get_seq_length(self, idx): |
| return self.seq_lengths[idx] |
|
|
| def get_all_actions(self): |
| return torch.cat([x[1] for x in self.data], dim=0) |
|
|
| def get_frames(self, idx, frames): |
| return [x[idx, frames] for x in self.data] |
|
|
| def __getitem__(self, idx): |
| return self.get_frames(idx, range(self.get_seq_length(idx))) |
|
|
| def __len__(self): |
| return len(self.seq_lengths) |
|
|
|
|
| def random_split_traj( |
| dataset: TrajectoryDataset, |
| lengths: Sequence[int], |
| generator: Optional[torch.Generator] = default_generator, |
| ) -> List[TrajectorySubset]: |
| """ |
| (Modified from torch.utils.data.dataset.random_split) |
| |
| Randomly split a trajectory dataset into non-overlapping new datasets of given lengths. |
| Optionally fix the generator for reproducible results, e.g.: |
| |
| >>> random_split_traj(range(10), [3, 7], generator=torch.Generator().manual_seed(42)) |
| |
| Args: |
| dataset (TrajectoryDataset): TrajectoryDataset to be split |
| lengths (sequence): lengths of splits to be produced |
| generator (Generator): Generator used for the random permutation. |
| """ |
| |
| if sum(lengths) != len(dataset): |
| raise ValueError( |
| "Sum of input lengths does not equal the length of the input dataset!" |
| ) |
|
|
| indices = randperm(sum(lengths), generator=generator).tolist() |
| return [ |
| TrajectorySubset(dataset, indices[offset - length : offset]) |
| for offset, length in zip(_accumulate(lengths), lengths) |
| ] |
|
|
|
|
| def split_traj_datasets(dataset, train_fraction=0.95, random_seed=42): |
| dataset_length = len(dataset) |
| lengths = [ |
| int(train_fraction * dataset_length), |
| dataset_length - int(train_fraction * dataset_length), |
| ] |
| train_set, val_set = random_split_traj( |
| dataset, lengths, generator=torch.Generator().manual_seed(random_seed) |
| ) |
| return train_set, val_set |
|
|
