| from collections import namedtuple |
| import os |
| import numpy as np |
| import torch |
| import tqdm |
|
|
| from .preprocessing import get_preprocess_fn |
| from .d4rl import load_environment, sequence_dataset |
| from .normalization import DatasetNormalizer |
| from .buffer import ReplayBuffer |
|
|
| Batch = namedtuple("Batch", "trajectories conditions") |
|
|
|
|
| class SequenceDataset(torch.utils.data.Dataset): |
|
|
| def __init__( |
| self, |
| env="avoiding-v0", |
| horizon=64, |
| normalizer="LimitsNormalizer", |
| preprocess_fns=[], |
| max_path_length=1000, |
| max_n_episodes=10000, |
| termination_penalty=0, |
| seed=None, |
| ): |
| self.preprocess_fn = get_preprocess_fn(preprocess_fns, env) |
| self.env_name = env |
| self.env = env = load_environment(env) |
| self.env.seed(seed) |
| self.horizon = horizon |
| self.max_path_length = max_path_length |
| itr = sequence_dataset(env, self.preprocess_fn) |
|
|
| fields = ReplayBuffer(max_n_episodes, max_path_length, termination_penalty) |
| for i, episode in enumerate(itr): |
| fields.add_path(episode) |
| fields.finalize() |
|
|
| normalizer_cache_path = os.path.join( |
| "logs", self.env_name, "normalizer", f"normalizer_H{horizon}.pth" |
| ) |
| print(f"[ datasets/sequence ] Calculating normalizer...") |
| self.normalizer = DatasetNormalizer( |
| fields, normalizer, path_lengths=fields["path_lengths"] |
| ) |
| os.makedirs(os.path.dirname(normalizer_cache_path), exist_ok=True) |
| torch.save(self.normalizer, normalizer_cache_path) |
| print(f"[ datasets/sequence ] Normalizer saved at {normalizer_cache_path}") |
|
|
| print("[ datasets/sequence ] Normalizer acquired") |
| self.indices = self.make_indices(fields.path_lengths, horizon) |
| print("[ datasets/sequence ] Indices made") |
|
|
| self.observation_dim = fields.observations.shape[-1] |
| self.action_dim = fields.actions.shape[-1] |
| self.fields = fields |
| self.n_episodes = fields.n_episodes |
| self.path_lengths = fields.path_lengths |
| self.normalize() |
| print("[ datasets/sequence ] Normalized") |
|
|
| print(fields) |
|
|
| def normalize(self, keys=["observations", "actions"]): |
| print("Normalizing...") |
| for key in tqdm.tqdm(keys): |
| array = self.fields[key].reshape(self.n_episodes * self.max_path_length, -1) |
| normed = self.normalizer(array, key) |
| self.fields[f"normed_{key}"] = normed.reshape( |
| self.n_episodes, self.max_path_length, -1 |
| ) |
|
|
| def make_indices(self, path_lengths, horizon): |
| indices = [] |
| for i, path_length in enumerate(path_lengths): |
| max_start = min(path_length - 1, self.max_path_length - horizon) |
| max_start = min(max_start, path_length - horizon) |
| for start in range(max_start): |
| end = start + horizon |
| indices.append((i, start, end)) |
| indices = np.array(indices) |
| return indices |
|
|
| def get_conditions(self, observations): |
| return {0: observations[0]} |
|
|
| def __len__(self): |
| return len(self.indices) |
|
|
| def __getitem__(self, idx, eps=1e-4): |
|
|
| path_ind, start, end = self.indices[idx] |
|
|
| observations = self.fields.normed_observations[path_ind, start:end] |
| actions = self.fields.normed_actions[path_ind, start:end] |
|
|
| conditions = self.get_conditions(observations) |
| trajectories = np.concatenate([actions, observations], axis=-1) |
| batch = Batch(trajectories, conditions) |
| return batch |
|
|
