From 8ce0aa14857a29dbd1a58b63929078978e490e51 Mon Sep 17 00:00:00 2001 From: allenzren Date: Sun, 8 Sep 2024 17:52:16 -0400 Subject: [PATCH] simplify pre-training dataset, use npz --- README.md | 5 +- agent/dataset/buffer.py | 110 --------------- agent/dataset/sequence.py | 128 ++++++------------ .../pretrain/avoid_m1/pre_diffusion_mlp.yaml | 4 +- .../pretrain/avoid_m1/pre_gaussian_mlp.yaml | 4 +- cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml | 4 +- .../pretrain/avoid_m2/pre_diffusion_mlp.yaml | 4 +- .../pretrain/avoid_m2/pre_gaussian_mlp.yaml | 4 +- cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml | 4 +- .../pretrain/avoid_m3/pre_diffusion_mlp.yaml | 4 +- .../pretrain/avoid_m3/pre_gaussian_mlp.yaml | 4 +- cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml | 4 +- .../pretrain/lamp_low/pre_diffusion_mlp.yaml | 4 +- .../pretrain/lamp_low/pre_diffusion_unet.yaml | 4 +- .../pretrain/lamp_low/pre_gaussian_mlp.yaml | 4 +- .../pretrain/lamp_med/pre_diffusion_mlp.yaml | 4 +- .../pretrain/lamp_med/pre_diffusion_unet.yaml | 4 +- .../pretrain/lamp_med/pre_gaussian_mlp.yaml | 4 +- .../one_leg_low/pre_diffusion_mlp.yaml | 4 +- .../one_leg_low/pre_diffusion_unet.yaml | 4 +- .../one_leg_low/pre_gaussian_mlp.yaml | 4 +- .../one_leg_med/pre_diffusion_mlp.yaml | 4 +- .../one_leg_med/pre_diffusion_unet.yaml | 4 +- .../one_leg_med/pre_gaussian_mlp.yaml | 4 +- .../round_table_low/pre_diffusion_mlp.yaml | 4 +- .../round_table_low/pre_diffusion_unet.yaml | 4 +- .../round_table_low/pre_gaussian_mlp.yaml | 4 +- .../round_table_med/pre_diffusion_mlp.yaml | 4 +- .../round_table_med/pre_diffusion_unet.yaml | 4 +- .../round_table_med/pre_gaussian_mlp.yaml | 4 +- .../pre_diffusion_mlp.yaml | 4 +- .../hopper-medium-v2/pre_diffusion_mlp.yaml | 4 +- .../walker2d-medium-v2/pre_diffusion_mlp.yaml | 4 +- .../pretrain/can/pre_diffusion_mlp.yaml | 4 +- .../pretrain/can/pre_diffusion_mlp_img.yaml | 4 +- .../pretrain/can/pre_diffusion_unet.yaml | 4 +- .../pretrain/can/pre_gaussian_mlp.yaml | 4 +- .../pretrain/can/pre_gaussian_mlp_img.yaml | 4 +- .../can/pre_gaussian_transformer.yaml | 4 +- cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml | 4 +- .../pretrain/can/pre_gmm_transformer.yaml | 4 +- .../pretrain/lift/pre_diffusion_mlp.yaml | 4 +- .../pretrain/lift/pre_diffusion_mlp_img.yaml | 4 +- .../pretrain/lift/pre_diffusion_unet.yaml | 4 +- .../pretrain/lift/pre_gaussian_mlp.yaml | 4 +- .../pretrain/lift/pre_gaussian_mlp_img.yaml | 4 +- .../lift/pre_gaussian_transformer.yaml | 4 +- cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml | 4 +- .../pretrain/lift/pre_gmm_transformer.yaml | 4 +- .../pretrain/square/pre_diffusion_mlp.yaml | 4 +- .../square/pre_diffusion_mlp_img.yaml | 4 +- .../pretrain/square/pre_diffusion_unet.yaml | 4 +- .../pretrain/square/pre_gaussian_mlp.yaml | 4 +- .../pretrain/square/pre_gaussian_mlp_img.yaml | 4 +- .../square/pre_gaussian_transformer.yaml | 4 +- .../pretrain/square/pre_gmm_mlp.yaml | 4 +- .../pretrain/square/pre_gmm_transformer.yaml | 4 +- .../pretrain/transport/pre_diffusion_mlp.yaml | 4 +- .../transport/pre_diffusion_mlp_img.yaml | 4 +- .../transport/pre_diffusion_unet.yaml | 4 +- .../pretrain/transport/pre_gaussian_mlp.yaml | 4 +- .../transport/pre_gaussian_mlp_img.yaml | 4 +- .../transport/pre_gaussian_transformer.yaml | 4 +- .../pretrain/transport/pre_gmm_mlp.yaml | 4 +- .../transport/pre_gmm_transformer.yaml | 4 +- model/diffusion/mlp_diffusion.py | 3 + 66 files changed, 170 insertions(+), 324 deletions(-) delete mode 100644 agent/dataset/buffer.py diff --git a/README.md b/README.md index c5a5d34..b31969b 100644 --- a/README.md +++ b/README.md @@ -57,7 +57,7 @@ export DPPO_DATA_DIR=/path/to/data --> Pre-training data for all tasks are pre-processed and can be found at [here](https://drive.google.com/drive/folders/1AXZvNQEKOrp0_jk1VLepKh_oHCg_9e3r?usp=drive_link). Pre-training script will download the data (including normalization statistics) automatically to the data directory. - + ### Run pre-training with data All the configs can be found under `cfg//pretrain/`. A new WandB project may be created based on `wandb.project` in the config file; set `wandb=null` in the command line to test without WandB logging. @@ -159,7 +159,8 @@ To use DDIM fine-tuning, set `denoising_steps=100` in pre-training and set `mode ## Adding your own dataset/environment ### Pre-training data -Pre-training script is at [`agent/pretrain/train_diffusion_agent.py`](agent/pretrain/train_diffusion_agent.py). The pre-training dataset [loader](agent/dataset/sequence.py) assumes a pickle file containing a dictionary of `observations`, `actions`, and `traj_length`, where `observations` and `actions` have the shape of num_episode x max_episode_length x obs_dim/act_dim, and `traj_length` is a 1-D array. One pre-processing example can be found at [`script/process_robomimic_dataset.py`](script/process_robomimic_dataset.py). +Pre-training script is at [`agent/pretrain/train_diffusion_agent.py`](agent/pretrain/train_diffusion_agent.py). The pre-training dataset [loader](agent/dataset/sequence.py) assumes a npz file containing numpy arrays `states`, `actions`, `images` (if using pixel) and `traj_length`, where `states` and `actions` have the shape of num_total_steps x obs_dim/act_dim, `images` num_total_steps x C (concatenated if multiple images) x H x W, and `traj_length` is a 1-D array for indexing across num_total_steps. + **Note:** The current implementation does not support loading history observations (only using observation at the current timestep). If needed, you can modify [here](agent/dataset/sequence.py#L130-L131). diff --git a/agent/dataset/buffer.py b/agent/dataset/buffer.py deleted file mode 100644 index 251c505..0000000 --- a/agent/dataset/buffer.py +++ /dev/null @@ -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() diff --git a/agent/dataset/sequence.py b/agent/dataset/sequence.py index 48b54dc..3a95fd3 100644 --- a/agent/dataset/sequence.py +++ b/agent/dataset/sequence.py @@ -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 diff --git a/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml b/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml index 056692f..649e4d3 100644 --- a/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m1/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: avoid_m1_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m1/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m1/train.npz seed: 42 device: cuda:0 @@ -63,7 +63,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml b/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml index 23e2884..c4fd533 100644 --- a/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m1/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: avoid_m1_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m1/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m1/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml b/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml index f20817f..0204457 100644 --- a/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m1/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: avoid_m1_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m1/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m1/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml b/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml index 482c7d3..6565320 100644 --- a/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m2/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: avoid_m2_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m2/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m2/train.npz seed: 42 device: cuda:0 @@ -63,7 +63,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml b/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml index d6c84a6..922d9d9 100644 --- a/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m2/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: avoid_m2_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m2/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m2/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml b/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml index 5f06db5..1b63a83 100644 --- a/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m2/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: avoid_m2_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m2/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m2/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml b/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml index e4d2864..9cba4ca 100644 --- a/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m3/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: avoid_m3_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m3/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m3/train.npz seed: 42 device: cuda:0 @@ -63,7 +63,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml b/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml index a1d1bfc..fa58f74 100644 --- a/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m3/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: avoid_m3_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m3/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m3/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml b/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml index f8dca2d..b577ef6 100644 --- a/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml +++ b/cfg/d3il/pretrain/avoid_m3/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: avoid_m3_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/d3il-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m3/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/d3il/avoid_m3/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml index cfef0c5..e22a8d2 100644 --- a/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml +++ b/cfg/furniture/pretrain/lamp_low/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -65,7 +65,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml index 5e50540..933a251 100644 --- a/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml +++ b/cfg/furniture/pretrain/lamp_low/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -67,7 +67,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml index caa91c9..2d700be 100644 --- a/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml +++ b/cfg/furniture/pretrain/lamp_low/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml index 646afc6..8d21455 100644 --- a/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml +++ b/cfg/furniture/pretrain/lamp_med/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -65,7 +65,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml index 1e26b27..7f87d43 100644 --- a/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml +++ b/cfg/furniture/pretrain/lamp_med/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -66,7 +66,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml index d65e4cb..8bed467 100644 --- a/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml +++ b/cfg/furniture/pretrain/lamp_med/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml index 49b8e59..80dce0c 100644 --- a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml +++ b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -65,7 +65,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml index 1b3e1fd..cccb378 100644 --- a/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml +++ b/cfg/furniture/pretrain/one_leg_low/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -67,7 +67,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml index 3357c4f..fe6da75 100644 --- a/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml +++ b/cfg/furniture/pretrain/one_leg_low/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml index b690ade..d8050cc 100644 --- a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml +++ b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -65,7 +65,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml index 0f5a3d5..22ddcfb 100644 --- a/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml +++ b/cfg/furniture/pretrain/one_leg_med/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -66,7 +66,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml index 05cdd7e..df64f22 100644 --- a/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml +++ b/cfg/furniture/pretrain/one_leg_med/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml index 6b6c9eb..6e32866 100644 --- a/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml +++ b/cfg/furniture/pretrain/round_table_low/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -65,7 +65,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml index 8186e2f..bd3c29b 100644 --- a/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml +++ b/cfg/furniture/pretrain/round_table_low/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -66,7 +66,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml index f9488a2..2d308d5 100644 --- a/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml +++ b/cfg/furniture/pretrain/round_table_low/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml b/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml index d854f3a..e7e3da9 100644 --- a/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml +++ b/cfg/furniture/pretrain/round_table_med/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -65,7 +65,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml b/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml index a3dc85a..cd85ebd 100644 --- a/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml +++ b/cfg/furniture/pretrain/round_table_med/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -66,7 +66,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml b/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml index 31764a5..543e16c 100644 --- a/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml +++ b/cfg/furniture/pretrain/round_table_med/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/furniture-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/furniture/${task}_${randomness}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml index bbc8772..6c83483 100644 --- a/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml +++ b/cfg/gym/pretrain/halfcheetah-medium-v2/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/gym-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/gym/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/gym/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml index 29894e5..236f13e 100644 --- a/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml +++ b/cfg/gym/pretrain/hopper-medium-v2/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/gym-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/gym/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/gym/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml b/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml index e638e9f..ab92bc0 100644 --- a/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml +++ b/cfg/gym/pretrain/walker2d-medium-v2/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/gym-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/gym/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/gym/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml index 7789a8f..8c9c588 100644 --- a/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml +++ b/cfg/robomimic/pretrain/can/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -61,7 +61,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml index 47faccd..5761803 100644 --- a/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml +++ b/cfg/robomimic/pretrain/can/pre_diffusion_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_img_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -82,7 +82,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml index 81e360c..bc73a00 100644 --- a/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml +++ b/cfg/robomimic/pretrain/can/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml index 269d08a..c887e17 100644 --- a/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml +++ b/cfg/robomimic/pretrain/can/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -53,7 +53,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml index a57e911..d14b255 100644 --- a/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml +++ b/cfg/robomimic/pretrain/can/pre_gaussian_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_img_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -74,7 +74,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml index 239d9d5..328dc18 100644 --- a/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml +++ b/cfg/robomimic/pretrain/can/pre_gaussian_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml index d9fd909..5ec7d9b 100644 --- a/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml +++ b/cfg/robomimic/pretrain/can/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml index 6cc1041..c8a7157 100644 --- a/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml +++ b/cfg/robomimic/pretrain/can/pre_gmm_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml index 5dc7b50..a0aad6a 100644 --- a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml +++ b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -61,7 +61,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml index c10c47b..2871525 100644 --- a/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml +++ b/cfg/robomimic/pretrain/lift/pre_diffusion_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_img_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -82,7 +82,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml index 09ca52b..f1872b1 100644 --- a/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml +++ b/cfg/robomimic/pretrain/lift/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml index 2d53bf4..b363574 100644 --- a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml +++ b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -53,7 +53,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml index 796d2ea..1f711b9 100644 --- a/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml +++ b/cfg/robomimic/pretrain/lift/pre_gaussian_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_img_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -74,7 +74,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml index f6a884b..f2c92e0 100644 --- a/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml +++ b/cfg/robomimic/pretrain/lift/pre_gaussian_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml index a4dbfaf..5941572 100644 --- a/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml +++ b/cfg/robomimic/pretrain/lift/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml index ef55a5f..f364b83 100644 --- a/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml +++ b/cfg/robomimic/pretrain/lift/pre_gmm_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml index 6d39d54..85540f8 100644 --- a/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml +++ b/cfg/robomimic/pretrain/square/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -62,7 +62,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml index 76bf985..7ca5e70 100644 --- a/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml +++ b/cfg/robomimic/pretrain/square/pre_diffusion_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_img_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -82,7 +82,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml index 2ca4271..cb3cac1 100644 --- a/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml +++ b/cfg/robomimic/pretrain/square/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml index 361271b..3e092af 100644 --- a/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml +++ b/cfg/robomimic/pretrain/square/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -53,7 +53,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml index 552a52f..9ae3035 100644 --- a/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml +++ b/cfg/robomimic/pretrain/square/pre_gaussian_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_img_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -74,7 +74,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml index 82c6a09..c116189 100644 --- a/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml +++ b/cfg/robomimic/pretrain/square/pre_gaussian_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml index b80301c..2efbaa1 100644 --- a/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml +++ b/cfg/robomimic/pretrain/square/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml index b3059dc..41c3fb4 100644 --- a/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml +++ b/cfg/robomimic/pretrain/square/pre_gmm_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml index 020c1af..6ecf14a 100644 --- a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml +++ b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -61,7 +61,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml index 08d2657..7731fcc 100644 --- a/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml +++ b/cfg/robomimic/pretrain/transport/pre_diffusion_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_mlp_img_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -83,7 +83,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml b/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml index 7b31552..55c691a 100644 --- a/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml +++ b/cfg/robomimic/pretrain/transport/pre_diffusion_unet.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_diffusion_agent.TrainDiffusionAgent name: ${env}_pre_diffusion_unet_ta${horizon_steps}_td${denoising_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -64,7 +64,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml index e3f8d86..8c68ec6 100644 --- a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml +++ b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -53,7 +53,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml index 9925dd9..9d475aa 100644 --- a/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml +++ b/cfg/robomimic/pretrain/transport/pre_gaussian_mlp_img.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_mlp_img_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}-img/train.npz seed: 42 device: cuda:0 @@ -75,7 +75,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset use_img: True dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml b/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml index ac79f9d..4beb459 100644 --- a/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml +++ b/cfg/robomimic/pretrain/transport/pre_gaussian_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gaussian_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml b/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml index 7e48026..13d1e24 100644 --- a/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml +++ b/cfg/robomimic/pretrain/transport/pre_gmm_mlp.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_mlp_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -55,7 +55,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml b/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml index 87c0f18..9255f10 100644 --- a/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml +++ b/cfg/robomimic/pretrain/transport/pre_gmm_transformer.yaml @@ -7,7 +7,7 @@ _target_: agent.pretrain.train_gaussian_agent.TrainGaussianAgent name: ${env}_pre_gmm_transformer_ta${horizon_steps} logdir: ${oc.env:DPPO_LOG_DIR}/robomimic-pretrain/${name}/${now:%Y-%m-%d}_${now:%H-%M-%S}_${seed} -train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.pkl +train_dataset_path: ${oc.env:DPPO_DATA_DIR}/robomimic/${env}/train.npz seed: 42 device: cuda:0 @@ -57,7 +57,7 @@ ema: decay: 0.995 train_dataset: - _target_: agent.dataset.sequence.StitchedActionSequenceDataset + _target_: agent.dataset.sequence.StitchedSequenceDataset dataset_path: ${train_dataset_path} horizon_steps: ${horizon_steps} cond_steps: ${cond_steps} diff --git a/model/diffusion/mlp_diffusion.py b/model/diffusion/mlp_diffusion.py index beccfde..aec8dec 100644 --- a/model/diffusion/mlp_diffusion.py +++ b/model/diffusion/mlp_diffusion.py @@ -124,6 +124,9 @@ class VisionDiffusionMLP(nn.Module): else: state = cond["state"] + # convert rgb to float32 for augmentation + rgb = rgb.float() + # get vit output - pass in two images separately if rgb.shape[1] == 6: # TODO: properly handle multiple images rgb1 = rgb[:, :3]