e2e_trainset = panda_datasets.PandaParticleFilterDataset(
"data/gentle_push_1000.hdf5",
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'),
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:
e2e_trainset = panda_datasets.PandaParticleFilterDataset(
"data/gentle_push_{}.hdf5".format(args.data_size),
subsequence_length=16,
particle_count=1,
particle_stddev=(.03, .03),
**dataset_args)
dataset_measurement = panda_datasets.PandaMeasurementDataset(
"data/gentle_push_{}.hdf5".format(args.data_size),
subsequence_length=16,