def preprocess_dataset()

in Experiments/Roboturk_DataLoader.py [0:0]

• 45 lines of code
• 6 McCabe index (conditional complexity)

	def preprocess_dataset(self):

		# for task_index in range(len(self.task_list)):
		# for task_index in [3,5]:
		for task_index in [0,1,2,4,6,7]:

			print("#######################################")
			print("Preprocessing task index: ", task_index)
			print("#######################################")

			# Get the name of environment.
			environment_name = self.files[task_index]['data'].attrs['env']
			# Create an actual robo-suite environment. 
			self.env = robosuite.make(environment_name)

			# Get sizes. 
			obs = self.env._get_observation()
			robot_state_size = obs['robot-state'].shape[0]
			object_state_size = obs['object-state'].shape[0]	

			# Create list of files for this task. 
			task_demo_list = []

			# For every element in the filelist of the element,
			for i in range(1,self.num_demos[task_index]+1):

				print("Preprocessing task index: ", task_index, " Demo Index: ", i, " of: ", self.num_demos[task_index])
			
				# Create list of datapoints for this demonstrations. 
				datapoint = {}

				# Get SEQUENCE of flattened states.
				try:
					flattened_state_sequence = self.files[task_index]['data/demo_{0}/states'.format(i)].value
					joint_action_sequence = self.files[task_index]['data/demo_{0}/joint_velocities'.format(i)].value
					gripper_action_sequence = self.files[task_index]['data/demo_{0}/gripper_actuations'.format(i)].value

					flattened_state_sequence = resample(flattened_state_sequence, flattened_state_sequence.shape[0]//self.ds_freq)

					number_timesteps = flattened_state_sequence.shape[0]
					robot_state_array = np.zeros((number_timesteps, robot_state_size))
					object_state_array = np.zeros((number_timesteps, object_state_size))

					# Get joint angle values from 
					joint_values = flattened_state_sequence[:,self.joint_angle_indices]
					# Get gripper values from state sequence. 
					gripper_finger_values = flattened_state_sequence[:,self.gripper_indices]

					# Normalize gripper values. 

					# 1 is right finger. 0 is left finger. 
					# 1-0 is right-left. 
		
					gripper_values = gripper_finger_values[:,1]-gripper_finger_values[:,0]
					gripper_values = (gripper_values-gripper_values.min()) / (gripper_values.max()-gripper_values.min())
					gripper_values = 2*gripper_values-1

					concatenated_demonstration = np.concatenate([joint_values,gripper_values.reshape((-1,1))],axis=1)
					concatenated_actions = np.concatenate([joint_action_sequence,gripper_action_sequence.reshape((-1,1))],axis=1)

					# For every element in sequence, set environment state. 
					for t in range(flattened_state_sequence.shape[0]):

						self.env.sim.set_state_from_flattened(flattened_state_sequence[t])

						# Now get observation. 
						observation = self.env._get_observation()

						# Robot and Object state appended to datapoint dictionary. 
						robot_state_array[t] = observation['robot-state']
						object_state_array[t] = observation['object-state']

				except: 

					datapoint['robot_state_array'] = np.zeros((1, robot_state_size))
					datapoint['object_state_array'] = np.zeros((1, object_state_size))				

				# Put both lists in a dictionary.
				datapoint['flat-state'] = flattened_state_sequence
				datapoint['robot-state'] = robot_state_array
				datapoint['object-state'] = object_state_array
				datapoint['demo'] = concatenated_demonstration				
				datapoint['demonstrated_actions'] = concatenated_actions

				# Add this dictionary to the file_demo_list. 
				task_demo_list.append(datapoint)

			# Create array.
			task_demo_array = np.array(task_demo_list)

			# Now save this file_demo_list. 
			np.save(os.path.join(self.dataset_directory,self.task_list[task_index],"New_Task_Demo_Array.npy"),task_demo_array)