<filename>datasets.py
import numpy as np
from sklearn import datasets as ds
import gym
import d4rl
import torch


class Dataset(torch.utils.data.Dataset):
    def __init__(self, data, labels=None):
        self.data = torch.tensor(data.astype(np.float32))
        self.labels = labels if labels is None else torch.tensor(labels.astype(np.float32))
        self.n_dims = self.data.shape[-1]

    def __getitem__(self, item):
        if self.labels is None:
            return self.data[item]
        return self.data[item], self.labels[item]

    def __len__(self):
        return len(self.data)


def get_datasets(name, conditional):
    if name == 'moons':
        assert not conditional
        x = ds.make_moons(n_samples=27000, shuffle=True, noise=0.05)[0]
        return Dataset(x[:24000]), Dataset(x[24000:])
    else:
        x = gym.make(name).get_dataset()
        done = np.logical_or(x['terminals'], x['timeouts'])
        trajs = np.split(x['actions'], np.nonzero(done)[0] + 1)
        # maze2d expert always gets 0 reward, never terminal; mujoco has no goal
        # terminal is last state of trajectory
        if conditional:
            data = np.concatenate([t[1:] for t in trajs], 0)
            labels = np.concatenate([t[:-1] for t in trajs], 0)
            train_data, val_data = np.split(data, [int(.9 * len(data))])
            train_labels, val_labels = np.split(labels, [int(.9 * len(labels))])
            return Dataset(train_data, train_labels), Dataset(val_data, val_labels)
        else:
            data = np.concatenate([np.concatenate([t[:-1], t[1:]], -1) for t in trajs], 0)
            train_data, val_data = np.split(data, [int(.9 * len(data))])
            return Dataset(train_data), Dataset(val_data)

<filename>datasets.py
import numpy as np
from sklearn import datasets as ds
import gym
import d4rl
import torch


class Dataset(torch.utils.data.Dataset):
    def __init__(self, data, labels=None):
        self.data = torch.tensor(data.astype(np.float32))
        self.labels = labels if labels is None else torch.tensor(labels.astype(np.float32))
        self.n_dims = self.data.shape[-1]

    def __getitem__(self, item):
        if self.labels is None:
            return self.data[item]
        return self.data[item], self.labels[item]

    def __len__(self):
        return len(self.data)


def get_datasets(name, conditional):
    if name == 'moons':
        assert not conditional
        x = ds.make_moons(n_samples=27000, shuffle=True, noise=0.05)[0]
        return Dataset(x[:24000]), Dataset(x[24000:])
    else:
        x = gym.make(name).get_dataset()
        done = np.logical_or(x['terminals'], x['timeouts'])
        trajs = np.split(x['actions'], np.nonzero(done)[0] + 1)
        # maze2d expert always gets 0 reward, never terminal; mujoco has no goal
        # terminal is last state of trajectory
        if conditional:
            data = np.concatenate([t[1:] for t in trajs], 0)
            labels = np.concatenate([t[:-1] for t in trajs], 0)
            train_data, val_data = np.split(data, [int(.9 * len(data))])
            train_labels, val_labels = np.split(labels, [int(.9 * len(labels))])
            return Dataset(train_data, train_labels), Dataset(val_data, val_labels)
        else:
            data = np.concatenate([np.concatenate([t[:-1], t[1:]], -1) for t in trajs], 0)
            train_data, val_data = np.split(data, [int(.9 * len(data))])
            return Dataset(train_data), Dataset(val_data)

# maze2d expert always gets 0 reward, never terminal; mujoco has no goal # terminal is last state of trajectory