| import pickle |
| from pathlib import Path |
|
|
| import einops |
| import torch |
| from datasets.core import TrajectoryDataset |
|
|
|
|
| class YorDataset(TrajectoryDataset): |
| """YOR bimanual episodes converted by utils/convert_yor.py. |
| |
| Layout under data_directory: |
| actions.pth FloatTensor [N, Tmax, A] (zero-padded) |
| seq_lengths.pkl list[int] |
| obses/episode_00000.pth uint8 Tensor [T, V, H, W, C] |
| meta.json action order / views / stride (informational) |
| |
| get_frames returns (obs, act, dummy_goal) with |
| obs: FloatTensor [T, V, C, H, W] in 0..1 |
| act: FloatTensor [T, A] |
| """ |
|
|
| def __init__(self, data_directory, prefetch: bool = True, views: int = 3): |
| self.data_directory = Path(data_directory) |
| self.actions = torch.load(self.data_directory / "actions.pth") |
| with open(self.data_directory / "seq_lengths.pkl", "rb") as f: |
| self.seq_lengths = pickle.load(f) |
| assert len(self.actions) == len(self.seq_lengths) |
| self.n = len(self.seq_lengths) |
| self.views = views |
| for i in range(self.n): |
| self.actions[i, self.seq_lengths[i]:] = 0 |
|
|
| self.prefetch = prefetch |
| self.obses = [] |
| if self.prefetch: |
| for i in range(self.n): |
| obs = torch.load(str(self.data_directory / "obses" / f"episode_{i:05d}.pth")) |
| self.obses.append(obs) |
|
|
| def get_seq_length(self, idx): |
| return self.seq_lengths[idx] |
|
|
| def get_all_actions(self): |
| result = [] |
| for i in range(self.n): |
| result.append(self.actions[i, :self.seq_lengths[i]]) |
| return torch.cat(result, dim=0) |
|
|
| def get_frames(self, idx, frames): |
| if self.prefetch: |
| obs = self.obses[idx][frames] |
| else: |
| obs = torch.load( |
| str(self.data_directory / "obses" / f"episode_{idx:05d}.pth")) |
| obs = obs[frames] |
| obs = einops.rearrange(obs, "T V H W C -> T V C H W")[:, : self.views] |
| obs = obs.float() / 255.0 |
| act = self.actions[idx, frames] |
| dummy_goal = torch.ones([obs.shape[0], 1, 1, 1]) |
| return obs, act, dummy_goal |
|
|
| def __getitem__(self, idx): |
| return self.get_frames(idx, range(self.get_seq_length(idx))) |
|
|
| def __len__(self): |
| return self.n |
|
|
