in Experiments/Roboturk_DataLoader.py [0:0]
def __init__(self, args):
super(Roboturk_NewSegmentedDataset, self).__init__()
self.dataset_directory = '/checkpoint/tanmayshankar/Roboturk/RoboTurkPilot'
self.args = args
# Require a task list.
# The task name is needed for setting the environment, rendering.
# We shouldn't need the environment for .. training though, should we?
self.task_list = ["bins-Bread", "bins-Can", "bins-Cereal", "bins-Milk", "pegs-RoundNut", "pegs-SquareNut"]
self.environment_names = ["SawyerPickPlaceBread","SawyerPickPlaceCan","SawyerPickPlaceCereal","SawyerPickPlaceMilk","SawyerNutAssemblyRound","SawyerNutAssemblySquare"]
self.num_demos = np.array([1069, 1069, 1069, 1069, 1144, 1145])
self.cummulative_num_demos = self.num_demos.cumsum()
self.cummulative_num_demos = np.insert(self.cummulative_num_demos,0,0)
# Append -1 to the start of cummulative_num_demos. This has two purposes.
# The first is that when we are at index 0 of the dataset, if we appended 0, np.searchsorted returns 0, rather than 1.
# For index 1, it returns 1. This was becoming inconsistent behavior for demonstrations in the same task.
# Now with -1 added to cumm_num_demos, when we are at task index 0, it would add -1 to the demo index. This is necessary for ALL tasks, not just the first...
# So that foils our really clever idea.
# Well, if the searchsorted returns the index of the equalling element, it probably consistently does this irrespective of vlaue.
# This means we can use this...
# No need for a clever solution, searchsorted has a "side" option that takes care of this.
self.total_length = self.num_demos.sum()
# Load data from all tasks.
self.files = []
# for i in range(len(self.task_list)):
for i in range(len(self.task_list)):
self.files.append( np.load("{0}/{1}/New_Task_Demo_Array.npy".format(self.dataset_directory, self.task_list[i]), allow_pickle=True))