def load_dataset(subsequence_length):
if args.dataset == "mujoco":
dataset = panda_datasets.PandaSubsequenceDataset(
"data/gentle_push_1000.hdf5",
# "data/gentle_push_1000.hdf5",
subsequence_length=subsequence_length,
**dataset_args)
elif args.dataset == "omnipush":
omnipush_train_files = (
"simpler/train0.hdf5",
"simpler/train1.hdf5",
"simpler/train2.hdf5",
"simpler/train3.hdf5",
"simpler/train4.hdf5",
"simpler/train5.hdf5",
)
dataset = omnipush_datasets.OmnipushSubsequenceDataset(
*omnipush_train_files,
subsequence_length=subsequence_length,
**dataset_args)
return dataset
subsequence_length=16,
particle_count=30,
particle_stddev=(.1, .1),
**dataset_args)
elif args.dataset == "omnipush":
omnipush_train_files = (
"simpler/train0.hdf5",
"simpler/train1.hdf5",
"simpler/train2.hdf5",
"simpler/train3.hdf5",
"simpler/train4.hdf5",
"simpler/train5.hdf5",
)
dynamics_trainset = omnipush_datasets.OmnipushDynamicsDataset(
*omnipush_train_files, **dataset_args)
dynamics_recurrent_trainset = omnipush_datasets.OmnipushSubsequenceDataset(
*omnipush_train_files, subsequence_length=16, **dataset_args)
measurement_trainset = omnipush_datasets.OmnipushMeasurementDataset(
*omnipush_train_files, samples_per_pair=10, **dataset_args)
e2e_trainset = omnipush_datasets.OmnipushParticleFilterDataset(
*omnipush_train_files,
subsequence_length=16,
particle_count=30,
particle_stddev=(.1, .1),
**dataset_args)
# Pre-train measurement
models = [
(pf_image_model, 'measurement_image'),
# (pf_force_model, 'measurement_force'),
]
measurement_trainset_loader = torch.utils.data.DataLoader(measurement_trainset,
"simpler/train0.hdf5",
"simpler/train1.hdf5",
"simpler/train2.hdf5",
"simpler/train3.hdf5",
"simpler/train4.hdf5",
"simpler/train5.hdf5",
subsequence_length=16,
stddev=(0.5, 0.5),
samples_per_pair=20,
**dataset_args)
dynamics_recurrent_trainset = omnipush_datasets.OmnipushSubsequenceDataset(
"simpler/train0.hdf5",
"simpler/train1.hdf5",
"simpler/train2.hdf5",
"simpler/train3.hdf5",
"simpler/train4.hdf5",
"simpler/train5.hdf5",
subsequence_length=32,
**dataset_args)
dataset_dynamics = omnipush_datasets.OmnipushDynamicsDataset(
"simpler/train0.hdf5",
"simpler/train1.hdf5",
"simpler/train2.hdf5",
"simpler/train3.hdf5",
"simpler/train4.hdf5",
"simpler/train5.hdf5",
subsequence_length=16,
**dataset_args)
else: