File size: 3,611 Bytes
7985e7e | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | 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
|