text: 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)
  