| import torch |
| import einops |
| import pickle |
| from pathlib import Path |
| from typing import Optional |
| from datasets.core import TrajectoryDataset |
|
|
|
|
| class PushTDataset(TrajectoryDataset): |
| def __init__( |
| self, |
| data_directory, |
| subset_fraction: Optional[float] = None, |
| relative=False, |
| ): |
| self.data_directory = Path(data_directory) |
| self.relative = relative |
| self.states = torch.load(self.data_directory / "states.pth") |
| if relative: |
| self.actions = torch.load(self.data_directory / "rel_actions.pth") |
| else: |
| self.actions = torch.load(self.data_directory / "abs_actions.pth") |
| with open(self.data_directory / "seq_lengths.pkl", "rb") as f: |
| self.seq_lengths = pickle.load(f) |
|
|
| 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.seq_lengths = self.seq_lengths[:n] |
|
|
| for i in range(n): |
| T = self.seq_lengths[i] |
| self.actions[i, T:] = 0 |
|
|
| def get_seq_length(self, idx): |
| return self.seq_lengths[idx] |
|
|
| def get_all_actions(self): |
| result = [] |
| for i in range(len(self.seq_lengths)): |
| T = self.seq_lengths[i] |
| result.append(self.actions[i, :T, :]) |
| return torch.cat(result, dim=0) |
|
|
| def get_frames(self, idx, frames): |
| vid_dir = self.data_directory / "obses" |
| obs = torch.load(str(vid_dir / f"episode_{idx:03d}.pth")) |
| obs = obs[frames] |
| obs = einops.rearrange(obs, "T H W C -> T 1 C H W") / 255.0 |
| act = self.actions[idx, frames] |
| mask = torch.ones(len(act)).bool() |
| return obs, act, mask |
|
|
| def __getitem__(self, idx): |
| return self.get_frames(idx, range(self.get_seq_length(idx))) |
|
|
| def __len__(self): |
| return len(self.seq_lengths) |
|
|
