File size: 613 Bytes
393d3de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import utils
import torch
import numpy as np
from pathlib import Path
from torch.utils.data import TensorDataset
from datasets.core import TrajectoryDataset
class YourTrajectoryDataset(TensorDataset, TrajectoryDataset):
def __init__(self, data_directory):
data_directory = Path(data_directory)
def get_seq_length(self, idx):
raise NotImplementedError
def get_frames(self, idx, frames):
raise NotImplementedError
# return obs / 255.0, actions, masks
def __getitem__(self, idx):
T = self.get_seq_length(idx)
return self.get_frames(idx, range(T))
|