dodox/dppo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

simplify pre-training dataset, use npz

Browse Source
This commit is contained in:
allenzren
2024-09-08 17:52:16 -04:00
parent 447c8dfd02
commit 8ce0aa1485
66 changed files with 170 additions and 324 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/dataset/buffer.py
-110
View File
@@ -1,110 +0,0 @@
"""
Pre-training data loader. Modified from https://github.com/jannerm/diffuser/blob/main/diffuser/datasets/buffer.py
"""
import numpy as np
import torch
def atleast_2d(x):
if isinstance(x, torch.Tensor):
while x.dim() < 2:
x = x.unsqueeze(-1)
return x
else:
while x.ndim < 2:
x = np.expand_dims(x, axis=-1)
return x
class StitchedBuffer:
def __init__(
self,
sum_of_path_lengths,
device="cpu",
):
self.sum_of_path_lengths = sum_of_path_lengths
if device == "cpu":
self._dict = {
"path_lengths": np.zeros(sum_of_path_lengths, dtype=int),
}
else:
self._dict = {
"path_lengths": torch.zeros(sum_of_path_lengths, dtype=int).to(device),
}
self._count = 0
self.sum_of_path_lengths = sum_of_path_lengths
self.device = device
def __repr__(self):
return "Fields:\n" + "\n".join(
f" {key}: {val.shape}" for key, val in self.items()
)
def __getitem__(self, key):
return self._dict[key]
def __setitem__(self, key, val):
self._dict[key] = val
self._add_attributes()
@property
def n_episodes(self):
return self._count
@property
def n_steps(self):
return sum(self["path_lengths"])
def _add_keys(self, path):
if hasattr(self, "keys"):
return
self.keys = list(path.keys())
def _add_attributes(self):
"""
can access fields with `buffer.observations`
instead of `buffer['observations']`
"""
for key, val in self._dict.items():
setattr(self, key, val)
def items(self):
return {k: v for k, v in self._dict.items() if k != "path_lengths"}.items()
def _allocate(self, key, array):
assert key not in self._dict
dim = array.shape[1:] # skip batch dimension
shape = (self.sum_of_path_lengths, *dim)
if self.device == "cpu":
self._dict[key] = np.zeros(shape, dtype=np.float32)
else:
self._dict[key] = torch.zeros(shape, dtype=torch.float32).to(self.device)
# print(f'[ utils/mujoco ] Allocated {key} with size {shape}')
def add_path(self, path):
path_length = len(path["observations"])
# assert path_length <= self.sum_of_path_lengths
## if first path added, set keys based on contents
self._add_keys(path)
## add tracked keys in path
for key in self.keys:
array = atleast_2d(path[key])
if key not in self._dict:
self._allocate(key, array)
self._dict[key][self._count : self._count + path_length] = array
## record path length
self._dict["path_lengths"][
self._count : self._count + path_length
] = path_length
## increment path counter
self._count += path_length
def finalize(self):
self._add_attributes()
agent/dataset/sequence.py
+40 -88
View File
@@ -8,7 +8,6 @@ No normalization is applied here --- we always normalize the data when pre-proce
"""
from collections import namedtuple
from tqdm import tqdm
import numpy as np
import torch
import logging
@@ -17,11 +16,7 @@ import random
log = logging.getLogger(__name__)
from .buffer import StitchedBuffer
Batch = namedtuple("Batch", "trajectories conditions")
ValueBatch = namedtuple("ValueBatch", "trajectories conditions values")
class StitchedSequenceDataset(torch.utils.data.Dataset):
@@ -32,7 +27,7 @@ class StitchedSequenceDataset(torch.utils.data.Dataset):
(tuple of) dimension of observation, action, images, etc.
Example:
Observations: [----------traj 1----------][---------traj 2----------] ... [---------traj N----------]
states: [----------traj 1----------][---------traj 2----------] ... [---------traj N----------]
Episode IDs: [---------- 1 ----------][---------- 2 ---------] ... [---------- N ---------]
"""
@@ -43,67 +38,56 @@ class StitchedSequenceDataset(torch.utils.data.Dataset):
cond_steps=1,
max_n_episodes=10000,
use_img=False,
device="cpu",
device="cuda:0",
):
self.horizon_steps = horizon_steps
self.cond_steps = cond_steps
self.device = device
self.use_img = use_img
# Load dataset to device specified
if dataset_path.endswith(".npz"):
dataset = np.load(dataset_path, allow_pickle=True)
dataset = np.load(dataset_path, allow_pickle=False) # only np arrays
else:
with open(dataset_path, "rb") as f:
dataset = pickle.load(f)
num_episodes = dataset["observations"].shape[0]
# Get the sum total of the valid trajectories' lengths
traj_lengths = dataset["traj_length"]
sum_of_path_lengths = np.sum(traj_lengths)
self.sum_of_path_lengths = sum_of_path_lengths
fields = StitchedBuffer(sum_of_path_lengths, device)
for i in tqdm(
range(min(max_n_episodes, num_episodes)), desc="Loading trajectories"
):
traj_length = traj_lengths[i]
episode = {
"observations": dataset["observations"][i][:traj_length],
"actions": dataset["actions"][i][:traj_length],
"episode_ids": i * np.ones(traj_length),
}
if use_img:
episode["images"] = dataset["images"][i][:traj_length]
for key, val in episode.items():
if device == "cpu":
episode[key] = val
else:
# if None array, save as empty tensor
if np.all(np.equal(episode[key], None)):
episode[key] = torch.empty(episode[key].shape).to(device)
else:
if key == "images":
episode[key] = torch.tensor(val, dtype=torch.uint8).to(
device
)
# (, H, W, C) -> (, C, H, W)
episode[key] = episode[key].permute(0, 3, 1, 2)
else:
episode[key] = torch.tensor(val, dtype=torch.float32).to(
device
)
fields.add_path(episode)
fields.finalize()
traj_lengths = dataset["traj_lengths"] # 1-D array
total_num_steps = np.sum(traj_lengths[:max_n_episodes])
# Set up indices for sampling
self.indices = self.make_indices(traj_lengths, horizon_steps)
self.obs_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.traj_lengths = traj_lengths
self.use_img = use_img
log.info(fields)
# Extract states and actions up to max_n_episodes
self.states = (
torch.from_numpy(dataset["states"][:total_num_steps]).float().to(device)
) # (total_num_steps, obs_dim)
self.actions = (
torch.from_numpy(dataset["actions"][:total_num_steps]).float().to(device)
) # (total_num_steps, action_dim)
log.info(f"Loaded dataset from {dataset_path}")
log.info(f"Number of episodes: {min(max_n_episodes, len(traj_lengths))}")
log.info(f"States shape/type: {self.states.shape, self.states.dtype}")
log.info(f"Actions shape/type: {self.actions.shape, self.actions.dtype}")
if self.use_img:
self.images = torch.from_numpy(dataset["images"][:total_num_steps]).to(
device
) # (total_num_steps, C, H, W)
log.info(f"Images shape/type: {self.images.shape, self.images.dtype}")
def __getitem__(self, idx):
start = self.indices[idx]
end = start + self.horizon_steps
states = self.states[start:end]
actions = self.actions[start:end]
if self.use_img:
images = self.images[start:end]
conditions = {
1 - self.cond_steps: {"state": states[0], "rgb": images[0]}
} # TODO: allow obs history, -1, -2, ...
else:
conditions = {1 - self.cond_steps: states[0]}
batch = Batch(actions, conditions)
return batch
def make_indices(self, traj_lengths, horizon_steps):
"""
@@ -119,44 +103,12 @@ class StitchedSequenceDataset(torch.utils.data.Dataset):
return indices
def set_train_val_split(self, train_split):
"""Not doing validation right now"""
num_train = int(len(self.indices) * train_split)
train_indices = random.sample(self.indices, num_train)
val_indices = [i for i in range(len(self.indices)) if i not in train_indices]
self.indices = train_indices
return val_indices
def set_indices(self, indices):
self.indices = indices
def get_conditions(self, observations, images=None):
"""
condition on current observation for planning. Take into account the number of conditioning steps.
"""
if images is not None:
return {
1 - self.cond_steps: {"state": observations[0], "rgb": images[0]}
} # TODO: allow obs history, -1, -2, ...
else:
return {1 - self.cond_steps: observations[0]}
def __len__(self):
return len(self.indices)
def __getitem__(self, idx, eps=1e-4):
raise NotImplementedError("Get item defined in subclass.")
class StitchedActionSequenceDataset(StitchedSequenceDataset):
"""Only use action trajectory, and then obs_cond for current observation"""
def __getitem__(self, idx):
start = self.indices[idx]
end = start + self.horizon_steps
observations = self.fields.observations[start:end]
actions = self.fields.actions[start:end]
images = None
if self.use_img:
images = self.fields.images[start:end]
conditions = self.get_conditions(observations, images)
batch = Batch(actions, conditions)
return batch