def visualize_robot_data()

in Experiments/PolicyManagers.py [0:0]

• 44 lines of code
• 5 McCabe index (conditional complexity)

	def visualize_robot_data(self):

		self.N = 100
		self.rollout_timesteps = self.args.traj_length
	
		if self.args.data=='MIME':
			self.visualizer = BaxterVisualizer()
			# self.state_dim = 16
		elif self.args.data=='Roboturk' or self.args.data=='OrigRoboturk' or self.args.data=='FullRoboturk':
			self.visualizer = SawyerVisualizer()
			# self.state_dim = 8
		elif self.args.data=='Mocap':
			self.visualizer = MocapVisualizer(args=self.args)
			# Because there are just more invalid DP's in Mocap.
			self.N = 100
		else: 
			self.visualizer = ToyDataVisualizer()

		self.latent_z_set = np.zeros((self.N,self.latent_z_dimensionality))		
		# These are lists because they're variable length individually.
		self.indices = []
		self.trajectory_set = []
		self.trajectory_rollout_set = []		

		model_epoch = int(os.path.split(self.args.model)[1].lstrip("Model_epoch"))

		self.rollout_gif_list = []
		self.gt_gif_list = []

		# Create save directory:
		upper_dir_name = os.path.join(self.args.logdir,self.args.name,"MEval")

		if not(os.path.isdir(upper_dir_name)):
			os.mkdir(upper_dir_name)

		self.dir_name = os.path.join(self.args.logdir,self.args.name,"MEval","m{0}".format(model_epoch))
		if not(os.path.isdir(self.dir_name)):
			os.mkdir(self.dir_name)

		self.max_len = 0

		for i in range(self.N):

			print("#########################################")	
			print("Getting visuals for trajectory: ",i)
			latent_z, sample_traj, sample_action_seq = self.run_iteration(0, i, return_z=True)

			if latent_z is not None:
				self.indices.append(i)

				if len(sample_traj)>self.max_len:
					self.max_len = len(sample_traj)

				self.latent_z_set[i] = copy.deepcopy(latent_z.detach().cpu().numpy())		
				
				trajectory_rollout = self.get_robot_visuals(i, latent_z, sample_traj)
				
				# self.trajectory_set[i] = copy.deepcopy(sample_traj)
				# self.trajectory_rollout_set[i] = copy.deepcopy(trajectory_rollout)	

				self.trajectory_set.append(copy.deepcopy(sample_traj))
				self.trajectory_rollout_set.append(copy.deepcopy(trajectory_rollout))

		# Get MIME embedding for rollout and GT trajectories, with same Z embedding. 
		embedded_z = self.get_robot_embedding()
		gt_animation_object = self.visualize_robot_embedding(embedded_z, gt=True)
		rollout_animation_object = self.visualize_robot_embedding(embedded_z, gt=False)

		self.write_embedding_HTML(gt_animation_object,prefix="GT")
		self.write_embedding_HTML(rollout_animation_object,prefix="Rollout")

		# Save webpage. 
		self.write_results_HTML()