| import os |
| import torch |
| import einops |
| import numpy as np |
| from pathlib import Path |
| from typing import Optional |
| from datasets.core import TrajectoryDataset |
|
|
|
|
| class PushMultiviewTrajectoryDataset(TrajectoryDataset): |
| def __init__( |
| self, |
| data_directory: os.PathLike, |
| onehot_goals=False, |
| subset_fraction: Optional[float] = None, |
| prefetch: bool = True, |
| view_idx: Optional[int] = None, |
| ): |
| self.data_directory = Path(data_directory) |
| self.states = np.load(self.data_directory / "multimodal_push_observations.npy") |
| self.actions = np.load(self.data_directory / "multimodal_push_actions.npy") |
| self.masks = np.load(self.data_directory / "multimodal_push_masks.npy") |
| self.view_idx = view_idx |
| self.subset_fraction = subset_fraction |
| if self.subset_fraction: |
| assert self.subset_fraction > 0 and self.subset_fraction <= 1 |
| n = int(len(self.states) * self.subset_fraction) |
| else: |
| n = len(self.states) |
| self.states = self.states[:n] |
| self.actions = self.actions[:n] |
| self.masks = self.masks[:n] |
|
|
| self.states = torch.from_numpy(self.states).float() |
| self.actions = torch.from_numpy(self.actions).float() / 0.03 |
| self.masks = torch.from_numpy(self.masks).bool() |
| self.prefetch = prefetch |
| if self.prefetch: |
| self.obses = [] |
| for i in range(n): |
| vid_path = self.data_directory / "obs_multiview" / f"{i:03d}.pth" |
| self.obses.append(torch.load(vid_path)) |
| self.onehot_goals = onehot_goals |
| if self.onehot_goals: |
| self.goals = torch.load(self.data_directory / "onehot_goals.pth").float() |
| self.goals = self.goals[:n] |
|
|
| def get_seq_length(self, idx): |
| return int(self.masks[idx].sum().item()) |
|
|
| def get_all_actions(self): |
| result = [] |
| |
| for i in range(len(self.masks)): |
| T = int(self.masks[i].sum().item()) |
| result.append(self.actions[i, :T, :]) |
| return torch.cat(result, dim=0) |
|
|
| def get_frames(self, idx, frames): |
| if self.prefetch: |
| obs = self.obses[idx][frames] |
| else: |
| obs = torch.load(self.data_directory / "obs_multiview" / f"{idx:03d}.pth")[ |
| frames |
| ] |
| obs = einops.rearrange(obs, "T V H W C -> T V C H W") / 255.0 |
| act = self.actions[idx, frames] |
| mask = self.masks[idx, frames] |
| if self.onehot_goals: |
| goal = self.goals[idx, frames] |
| |
| return obs, act, goal |
| else: |
| dummy_goal = torch.ones([obs.shape[0], 1, 1, 1]) |
| return obs, act, dummy_goal |
|
|
| def __getitem__(self, idx): |
| T = self.masks[idx].sum().int().item() |
| output = self.get_frames(idx, range(T)) |
| if self.view_idx is not None: |
| output = ( |
| output[0][:, self.view_idx : self.view_idx + 1], |
| output[1], |
| output[2], |
| ) |
| return output |
|
|
| def __len__(self): |
| return len(self.states) |
|
|
