File size: 2,711 Bytes
a3f6fc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import torch
import einops
import pickle
from pathlib import Path
from typing import Optional
from datasets.core import TrajectoryDataset
# import numpy as np
# from torch.serialization import safe_globals
class CubeDataset(TrajectoryDataset):
def __init__(
self,
data_directory,
subset_fraction: Optional[float] = None,
prefetch: bool = True,
*args,
**kwargs,
):
self.data_directory = Path(data_directory)
self.states = torch.load(self.data_directory / "latents.pth")
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)
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 # redo zero padding
self.prefetch = prefetch
if self.prefetch:
self.obses = []
for i in range(n):
vid_dir = self.data_directory / "obses"
obs = torch.load(str(vid_dir / f"episode_{i:05d}.pth"), weights_only=False)
self.obses.append(obs) # THWC
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):
if self.prefetch:
obs = self.obses[idx][frames]
else:
vid_dir = self.data_directory / "obses"
obs = torch.load(
str(vid_dir / f"episode_{idx:05d}.pth"), weights_only=False
)
obs = obs[frames]
obs = einops.rearrange(obs, "T H W C -> T 1 C H W") / 255.0 # T V C H W, 1 view
obs = torch.from_numpy(obs).float()
act = self.actions[idx, frames]
mask = torch.ones(len(act)).bool()
dummy_goal = torch.ones([obs.shape[0], 1, 1, 1]) # dummy goal, T V P E
# return obs, act, mask, goal
return obs, act, dummy_goal
def __getitem__(self, idx):
return self.get_frames(idx, range(self.get_seq_length(idx)))
def __len__(self):
return len(self.seq_lengths)
|