def get_image_list_for_object_traj(output_rotation, output_position, target_rotation, target_position, gt_contact_point, initial_rotation, initial_position, object_name, environment, input_rgb, output_contact_point=None): env_state = EnvState(position=initial_position, rotation=initial_rotation, object_name=object_name) initial_image = environment.get_rgb_for_position_rotation(object_state=env_state) initial_double_image = np.stack([initial_image, initial_image], axis=0) # because we need it both for target and output all_images = [initial_double_image] for seq_ind in range(output_rotation.shape[0]): output_env_state = EnvState(position=output_position[seq_ind], rotation=output_rotation[seq_ind], object_name=object_name) output_image = environment.get_rgb_for_position_rotation(object_state=output_env_state, contact_point=output_contact_point) # add force viz, forces_poc=force_visualization[seq_ind]) target_env_state = EnvState(position=target_position[seq_ind], rotation=target_rotation[seq_ind], object_name=object_name) target_image = environment.get_rgb_for_position_rotation(object_state=target_env_state, contact_point=gt_contact_point) # combine images combined_images = np.stack([target_image, output_image], axis=0) all_images.append(combined_images) # stack images on top, base on numpy or pytorch. it has to be of shape # batch, time, channels, height, width, image_list = np.stack(all_images, axis=1) # convert to tensor and make it channel first, we have to convert it to [0,1] because tensorboardX does not work with 0-255 image_list = torch.Tensor(image_list).permute(0, 1, 4, 2, 3).float() / 255. image_shapes = (image_list.shape[-2], image_list.shape[-1]) # adding rgb images rgb_image = (F.interpolate(normalize(input_rgb), size=image_shapes)) rgb_image = rgb_image.unsqueeze(0) image_list = torch.cat([rgb_image, image_list], dim=0) image_list = image_list.permute(1, 0, 3, 4, 2) image_list = add_border_to_images(image_list) return image_list
def get_image_list_for_keypoints(full_target_keypoints, full_output_keypoints, full_output_position, full_output_rotation, gt_contact_point, rgb_image, object_name, environment): # adding rgb images rgb_image = (F.interpolate(normalize(rgb_image), size=(environment.qualitative_size, environment.qualitative_size))).cpu().detach() rgb_image = rgb_image.unsqueeze(0) rgb_image_channel_last = rgb_image.squeeze(0).permute(0, 2, 3, 1) all_images = [] for seq_ind in range(full_output_position.shape[0]): output_env_state = EnvState(position=full_output_position[seq_ind], rotation=full_output_rotation[seq_ind], object_name=object_name) output_image = environment.get_rgb_for_position_rotation(object_state=output_env_state, contact_point=gt_contact_point) # add force viz, forces_poc=force_visualization[seq_ind]) output_image = torch.Tensor(output_image).float() / 255. image_for_target_keypoints = put_keypoints_on_image(rgb_image_channel_last[seq_ind], full_target_keypoints[seq_ind]) image_for_output_keypoints = put_keypoints_on_image(rgb_image_channel_last[seq_ind], full_output_keypoints[seq_ind]) combined_images = torch.stack([image_for_target_keypoints, image_for_output_keypoints, output_image], dim=0) all_images.append(combined_images) # stack images on top, base on numpy or pytorch. it has to be of shape # batch, time, channels, height, width, image_list = torch.stack(all_images, dim=1) # convert to tensor and make it channel first, we have to convert it to [0,1] because tensorboardX does not work with 0-255 image_list = image_list.permute(0, 1, 4, 2, 3).float() image_list = torch.cat([rgb_image, image_list], dim=0) image_list = image_list.permute(1, 0, 3, 4, 2) image_list = add_border_to_images(image_list) return image_list
def get_image_list_for_viz_cp(output_cp, target_cp, full_target_position, full_target_rotation, object_name, input_rgb, environment): all_images = [] for seq_ind in range(output_cp.shape[0]): env_state = EnvState(position=full_target_position[seq_ind], rotation=full_target_rotation[seq_ind], object_name=object_name) target_image = environment.get_rgb_for_position_rotation(object_state=env_state, contact_point=target_cp, contact_points_size=1) output_image = environment.get_rgb_for_position_rotation(object_state=env_state, contact_point=output_cp[seq_ind], contact_points_size=1) # combine images combined_images = np.stack([target_image, output_image], axis=0) all_images.append(combined_images) # stack images on top, base on numpy or pytorch. it has to be of shape # batch, time, channels, height, width, image_list = np.stack(all_images, axis=1) # convert to tensor and make it channel first, we have to convert it to [0,1] because tensorboardX does not work with 0-255 image_list = torch.Tensor(image_list).permute(0, 1, 4, 2, 3).float() / 255. image_shapes = (image_list.shape[-2], image_list.shape[-1]) # adding rgb images rgb_image = (F.interpolate(normalize(input_rgb), size=image_shapes)) rgb_image = rgb_image.unsqueeze(0) image_list = torch.cat([rgb_image, image_list], dim=0) image_list = image_list.permute(1, 0, 3, 4, 2) image_list = add_border_to_images(image_list) return image_list
def get_current_state(self, object_num=None): if object_num is None: object_num = self.object_of_interest_id objectPos, objectOr = self._p.getBasePositionAndOrientation(object_num) objectOr = quaternion_bullet2normal(objectOr) velocity, omega = self._p.getBaseVelocity(bodyUniqueId=object_num) return EnvState(object_name=self.object_name, position=objectPos, rotation=objectOr, velocity=velocity, omega=omega)
def forward(self, input_dict, target): initial_position = input_dict['initial_position'] initial_rotation = input_dict['initial_rotation'] rgb = input_dict['rgb'] batch_size, seq_len, c, w, h = rgb.shape object_name = input_dict['object_name'] assert len(object_name) == 1 # only support one object object_name = object_name[0] image_features = self.resnet_features(rgb) # Contact point prediction tower image_features_contact_point = self.contact_point_image_embed( image_features.view(batch_size * seq_len, 512, 7, 7)).view(batch_size, seq_len, 64 * 7 * 7) initial_object_features_contact_point = \ self.contact_point_input_object_embed(torch.cat([initial_position, initial_rotation], dim=-1)) object_features_contact_point = initial_object_features_contact_point.unsqueeze(1).\ repeat(1, self.sequence_length, 1) # add a dimension for sequence length and then repeat that # Predict contact point input_embedded_contact_point = torch.cat( [image_features_contact_point, object_features_contact_point], dim=-1) embedded_sequence_contact_point, ( _, _) = self.contact_point_encoder(input_embedded_contact_point) contact_points_prediction = self.contact_point_decoder( embedded_sequence_contact_point).view(batch_size, seq_len, self.number_of_cp, 3)[:, -1, :, :] # Predict contact point for each image and get the last prediction # Force prediction tower image_features_force = self.image_embed( image_features.view(batch_size * seq_len, 512, 7, 7)).view(batch_size, seq_len, 64 * 7 * 7) initial_object_features_force = self.input_object_embed( torch.cat([initial_position, initial_rotation], dim=-1)) object_features_force = initial_object_features_force.unsqueeze( 1).repeat(1, self.sequence_length, 1) # add a dimension for sequence length and then repeat that input_embedded_force = torch.cat( [image_features_force, object_features_force], dim=-1) embedded_sequence_force, ( hidden_force, cell_force) = self.lstm_encoder(input_embedded_force) last_hidden = hidden_force.view(self.num_layers, 1, 1, self.hidden_size)[-1, -1, :, :] # num_layers, num direction, batchsize, hidden and then take the last hidden layer last_cell = cell_force.view(self.num_layers, 1, 1, self.hidden_size)[-1, -1, :, :] # num_layers, num direction, batchsize, cell and then take the last hidden layer hn = last_hidden cn = last_cell resulting_force_success = [] forces_directions = [] all_force_applied = [] env_state = EnvState(object_name=object_name, rotation=initial_rotation[0], position=initial_position[0], velocity=None, omega=None) resulting_position = [] resulting_rotation = [] for seq_ind in range(self.sequence_length - 1): prev_location = env_state.toTensorCoverName().unsqueeze(0) contact_point_as_input = contact_points_prediction.view(1, 3 * 5) prev_state_and_cp = torch.cat( [prev_location, contact_point_as_input], dim=-1) prev_state_embedded = self.state_embed(prev_state_and_cp) next_state_embedded = embedded_sequence_force[:, seq_ind + 1] input_lstm_cell = torch.cat( [prev_state_embedded, next_state_embedded], dim=-1) (hn, cn) = self.lstm_decoder(input_lstm_cell, (hn, cn)) force = self.forces_directions_decoder(hn) assert force.shape[0] == 1 force = force.squeeze(0) assert force.shape[0] == (self.number_of_cp * 3) # initial initial_velocity is whatever it was the last frame, # note that the gradients are not backproped here env_state, force_success, force_applied = \ self.environment_layer.apply(self.environment, env_state.toTensor(), force, contact_point_as_input.squeeze(0)) env_state = EnvState.fromTensor(env_state) resulting_position.append(env_state.position) resulting_rotation.append(env_state.rotation) resulting_force_success.append(force_success) forces_directions.append(force.view(self.number_of_cp, 3)) all_force_applied.append(force_applied) resulting_position = torch.stack(resulting_position, dim=0) resulting_rotation = torch.stack(resulting_rotation, dim=0) resulting_force_success = torch.stack(resulting_force_success, dim=0) forces_directions = torch.stack(forces_directions, dim=0) all_force_applied = torch.stack(all_force_applied, dim=0) resulting_position = resulting_position.unsqueeze( 0) # adding batchsize back because we need it in the loss resulting_rotation = resulting_rotation.unsqueeze( 0) # adding batchsize back because we need it in the loss forces_directions = forces_directions.unsqueeze(0) resulting_force_success = resulting_force_success.unsqueeze(0) all_force_applied = all_force_applied.unsqueeze(0) all_keypoints = get_keypoint_projection( object_name, resulting_position, resulting_rotation, self.all_objects_keypoint_tensor[object_name]) all_keypoints = all_keypoints.unsqueeze( 0) # adding batchsize back because we need it in the loss contact_points_prediction = contact_points_prediction.unsqueeze( 1).repeat(1, seq_len, 1, 1) output_dict = { 'keypoints': all_keypoints, 'rotation': resulting_rotation, 'position': resulting_position, 'force_success_flag': resulting_force_success, 'force_applied': all_force_applied, 'force_direction': forces_directions, # batch size x seq len -1 x number of cp x 3 'contact_points': contact_points_prediction, } target['object_name'] = input_dict['object_name'] return output_dict, target
def forward(self, input, target): initial_position = input['initial_position'] initial_rotation = input['initial_rotation'] contact_points = input['contact_points'] assert contact_points.shape[0] == 1 contact_points = contact_points.squeeze(0) object_name = input['object_name'] assert len(object_name) == 1 object_name = object_name[0] target['object_name'] = input['object_name'] predefined_force = target['forces'].squeeze(0).view(self.sequence_length - 1, 5 * 3) all_forces = torch.nn.Parameter(predefined_force.detach()) loss_function = self.this_loss_func if self.gpu_ids != -1: all_forces = all_forces.cuda().detach() loss_function = loss_function.cuda() optimizer = torch.optim.SGD([all_forces], lr=self.base_lr) number_of_states = 20 for t in range(number_of_states): # all_forces = all_forces.clamp(-1.5, 1.5) if t <= self.step_size: lr = self.base_lr elif t <= self.step_size * 2: lr = self.base_lr * 0.1 elif t <= self.step_size * 3: lr = self.base_lr * 0.01 for param_group in optimizer.param_groups: param_group['lr'] = lr env_state = EnvState(object_name=object_name, rotation=initial_rotation[0], position=initial_position[0], velocity=None, omega=None) resulting_force_success = [] forces_directions = [] all_force_applied = [] resulting_position = [] resulting_rotation = [] for seq_ind in range(self.sequence_length - 1): force = all_forces[seq_ind] assert force.shape[0] == (self.number_of_cp * 3) # initial initial_velocity is whatever it was the last frame, note that the gradients are not backproped here env_state, force_success, force_applied = self.environment_layer.apply(self.environment, env_state.toTensor(), force, contact_points) env_state = EnvState.fromTensor(env_state) resulting_position.append(env_state.position) resulting_rotation.append(env_state.rotation) resulting_force_success.append(force_success) forces_directions.append(force.view(self.number_of_cp, 3)) all_force_applied.append(force_applied) resulting_position = torch.stack(resulting_position, dim=0) resulting_rotation = torch.stack(resulting_rotation, dim=0) resulting_position = resulting_position.unsqueeze(0) # adding batchsize back because we need it in the loss resulting_rotation = resulting_rotation.unsqueeze(0) # adding batchsize back because we need it in the loss all_keypoints = get_keypoint_projection(object_name, resulting_position, resulting_rotation, self.all_objects_keypoint_tensor[object_name]) all_keypoints = all_keypoints.unsqueeze(0) # adding batchsize back because we need it in the loss output = { 'keypoints': all_keypoints, } loss = loss_function(output, target) loss.backward() optimizer.step() optimizer.zero_grad() resulting_force_success = torch.stack(resulting_force_success, dim=0) forces_directions = torch.stack(forces_directions, dim=0) all_force_applied = torch.stack(all_force_applied, dim=0) forces_directions = forces_directions.unsqueeze(0) resulting_force_success = resulting_force_success.unsqueeze(0) all_force_applied = all_force_applied.unsqueeze(0) all_keypoints = torch.tensor(all_keypoints, requires_grad=True) output = { 'keypoints': all_keypoints, 'rotation': resulting_rotation, 'position': resulting_position, 'force_success_flag': resulting_force_success, 'force_applied': all_force_applied, 'force_direction': forces_directions, # batch size x seq len -1 x number of cp x 3 } return output, target
def forward(self, input, target): object_name = input['object_name'] assert len(object_name) == 1 object_name = object_name[0] initial_keypoint = input['initial_keypoint'] initial_keypoint = initial_keypoint / DEFAULT_IMAGE_SIZE initial_keypoint = initial_keypoint.view(1, 2 * 10) object_points = self.all_objects_keypoint_tensor[object_name].view( 1, 3 * 10) rgb = input['rgb'] contact_points = input['contact_points'] assert contact_points.shape[0] == 1 contact_points = contact_points.squeeze(0) contact_point_as_input = input['contact_points'].view( 1, 3 * 5) # Batchsize , 3 * number of contact points batch_size, seq_len, c, w, h = rgb.shape image_features = self.resnet_features(rgb) image_features = self.image_embed( image_features.view(batch_size * seq_len, 512, 7, 7)).view(batch_size, seq_len, 64 * 7 * 7) initial_state = self.predict_initial_pose( torch.cat([initial_keypoint, object_points], dim=-1)).view(1, 3 + 4) initial_position = initial_state[:, :3] initial_rotation = initial_state[:, 3:] initial_rotation = F.normalize(initial_rotation, dim=-1) # Just for visualization purposes input['initial_rotation'] = initial_rotation input['initial_position'] = initial_position initial_object_features = self.input_object_embed( torch.cat([initial_position, initial_rotation], dim=-1)) object_features = initial_object_features.unsqueeze(1).repeat( 1, self.sequence_length, 1) # add a dimension for sequence length and then repeat that input_embedded = torch.cat([image_features, object_features], dim=-1) embedded_sequence, (hidden, cell) = self.lstm_encoder(input_embedded) last_hidden = hidden.view( self.num_layers, 1, 1, self.hidden_size )[-1, -1, :, :] # num_layers, num direction, batchsize, hidden and then take the last hidden layer last_cell = cell.view( self.num_layers, 1, 1, self.hidden_size )[-1, -1, :, :] # num_layers, num direction, batchsize, cell and then take the last hidden layer hn = last_hidden cn = last_cell resulting_force_success = [] forces_directions = [] all_force_applied = [] env_state = EnvState(object_name=object_name, rotation=initial_rotation[0], position=initial_position[0], velocity=None, omega=None) resulting_position = [] resulting_rotation = [] for seq_ind in range(self.sequence_length - 1): prev_location = env_state.toTensorCoverName().unsqueeze(0) prev_state_and_cp = torch.cat( [prev_location, contact_point_as_input], dim=-1) prev_state_embedded = self.state_embed(prev_state_and_cp) next_state_embedded = embedded_sequence[:, seq_ind + 1] input_lstm_cell = torch.cat( [prev_state_embedded, next_state_embedded], dim=-1) (hn, cn) = self.lstm_decoder(input_lstm_cell, (hn, cn)) force = self.forces_directions_decoder(hn) assert force.shape[0] == 1 force = force.squeeze(0) assert force.shape[0] == (self.number_of_cp * 3) # Cap forces # think about this more force = force.clamp(-1., 1.) # initial initial_velocity is whatever it was the last frame, note that the gradients are not backproped here env_state, force_success, force_applied = self.environment_layer.apply( self.environment, env_state.toTensor(), force, contact_points) env_state = EnvState.fromTensor(env_state) resulting_position.append(env_state.position) resulting_rotation.append(env_state.rotation) resulting_force_success.append(force_success) forces_directions.append(force.view(self.number_of_cp, 3)) all_force_applied.append(force_applied) resulting_position = torch.stack(resulting_position, dim=0) resulting_rotation = torch.stack(resulting_rotation, dim=0) resulting_force_success = torch.stack(resulting_force_success, dim=0) forces_directions = torch.stack(forces_directions, dim=0) all_force_applied = torch.stack(all_force_applied, dim=0) resulting_position = resulting_position.unsqueeze( 0) # adding batchsize back because we need it in the loss resulting_rotation = resulting_rotation.unsqueeze( 0) # adding batchsize back because we need it in the loss forces_directions = forces_directions.unsqueeze(0) resulting_force_success = resulting_force_success.unsqueeze(0) all_force_applied = all_force_applied.unsqueeze(0) all_keypoints = get_keypoint_projection( object_name, resulting_position, resulting_rotation, self.all_objects_keypoint_tensor[object_name]) all_keypoints = all_keypoints.unsqueeze( 0) # adding batchsize back because we need it in the loss output = { 'keypoints': all_keypoints, 'rotation': resulting_rotation, 'position': resulting_position, 'force_success_flag': resulting_force_success, 'force_applied': all_force_applied, 'force_direction': forces_directions, # batch size x seq len -1 x number of cp x 3 } target['object_name'] = input['object_name'] return output, target
def forward(self, input, target): initial_position = input['initial_position'] initial_rotation = input['initial_rotation'] rgb = input['rgb'] batch_size, seq_len, c, w, h = rgb.shape contact_point_as_input = input['contact_points'].view( batch_size, 5 * 3) image_features = self.resnet_features(rgb) image_features = self.image_embed( image_features.view(batch_size * seq_len, 512, 7, 7)).view(batch_size, seq_len, 64 * 7 * 7) initial_object_features = self.input_object_embed( torch.cat([initial_position, initial_rotation], dim=-1)) object_features = initial_object_features.unsqueeze(1).repeat( 1, self.sequence_length, 1) # add a dimension for sequence length and then repeat that input_embedded = torch.cat([image_features, object_features], dim=-1) embedded_sequence, (hidden, cell) = self.lstm_encoder(input_embedded) contact_point_embedding = self.contact_point_embed( contact_point_as_input).unsqueeze(1).repeat(1, seq_len, 1) combined_w_cp = torch.cat([embedded_sequence, contact_point_embedding], dim=-1) forces_prediction = self.force_decoder(combined_w_cp).view( batch_size, seq_len, self.number_of_cp, 3) # Predict contact point for each image forces_prediction = forces_prediction[:, 1:, :, :] forces_prediction = forces_prediction.clamp(-1.5, 1.5) output = { 'forces': forces_prediction, # batchsize x seq len x number of cp x 3 } target['object_name'] = input['object_name'] # forces_prediction[:] = 0. # output['forces'][:] = 0. # remove # forces_prediction[:] = target['forces'] if not self.train_mode: object_name = input['object_name'] assert len(object_name) == 1 object_name = object_name[0] contact_points = input['contact_points'] assert contact_points.shape[0] == 1 contact_points = contact_points.squeeze(0) resulting_position = [] resulting_rotation = [] env_state = EnvState(object_name=object_name, rotation=initial_rotation[0], position=initial_position[0], velocity=None, omega=None) for seq_ind in range(self.sequence_length - 1): force = forces_prediction[0, seq_ind].view(self.number_of_cp * 3) # initial initial_velocity is whatever it was the last frame, note that the gradients are not backproped here env_state, force_success, force_applied = self.environment_layer.apply( self.environment, env_state.toTensor(), force, contact_points) env_state = EnvState.fromTensor(env_state) resulting_position.append(env_state.position) resulting_rotation.append(env_state.rotation) resulting_position = torch.stack(resulting_position, dim=0) resulting_rotation = torch.stack(resulting_rotation, dim=0) resulting_position = resulting_position.unsqueeze( 0) # adding batchsize back because we need it in the loss resulting_rotation = resulting_rotation.unsqueeze( 0) # adding batchsize back because we need it in the loss all_keypoints = get_keypoint_projection( object_name, resulting_position, resulting_rotation, self.all_objects_keypoint_tensor[object_name]) all_keypoints = all_keypoints.unsqueeze( 0) # adding batchsize back because we need it in the loss output['keypoints'] = all_keypoints output['rotation'] = resulting_rotation output['position'] = resulting_position output['force_applied'] = output['forces'] output['force_direction'] = output['forces'] return output, target
def forward(ctx, environment, initial_state, forces_tensor, list_of_contact_points): assert len(forces_tensor.shape) == 1, 'The force should be flattened' initial_state = EnvState.fromTensor(initial_state) forces = ForceValOnly.fromForceArray( forces_tensor.view(environment.number_of_cp, 3)) list_of_contact_points = list_of_contact_points.view(5, 3) finite_diff_success_flags = [] with torch.no_grad(): env_state, force_success_flag, force_that_applied = environment.init_location_and_apply_force( initial_state=initial_state, forces=forces, list_of_contact_points=list_of_contact_points) f_x_env_state_tensor = env_state.toTensor() no_grad_on = not (True in ctx.needs_input_grad) if not no_grad_on: # at least one of them needs gradients with torch.no_grad(): cp_h = 0.05 # position and linear velocity should not depend, force similar, just angular manual_tweaks_cp = [] contact_point_tensor = list_of_contact_points.view(-1) for arg_ind in range(len(contact_point_tensor)): changed_contact_point = contact_point_tensor * 1 # Just copy tweak_value_cp = cp_h if random.random() > 0.5: tweak_value_cp *= -1 changed_contact_point[arg_ind] += tweak_value_cp changed_contact_point = changed_contact_point.view(5, 3) env_state_h, _, _ = environment.init_location_and_apply_force( initial_state=initial_state, forces=forces, list_of_contact_points=changed_contact_point) manual_tweaks_cp.append( (env_state_h.toTensor() - f_x_env_state_tensor) / (tweak_value_cp)) ctx.contact_pt_x_plus_h_diff = torch.stack(manual_tweaks_cp, dim=-2) force_h = environment.force_h manual_tweaks_force = [] for arg_ind in range(len(forces_tensor)): changed_force = forces_tensor * 1 # Just copy tweak_value_force = force_h if random.random() > 0.5: tweak_value_force *= -1 changed_force[arg_ind] += tweak_value_force changed_force = ForceValOnly.fromForceArray( changed_force.view(environment.number_of_cp, -1)) env_state_h, force_success_flag_h, force_that_applied_h = environment.init_location_and_apply_force( initial_state=initial_state, forces=changed_force, list_of_contact_points=list_of_contact_points) manual_tweaks_force.append( (env_state_h.toTensor() - f_x_env_state_tensor) / (tweak_value_force)) finite_diff_success_flags.append(force_success_flag_h) ctx.f_x_plus_h_diff = torch.stack(manual_tweaks_force, dim=-2) state = initial_state.toTensor() manual_tweaks_initial_state = [] state_h = environment.state_h for arg_ind in range(len(state)): if arg_ind == EnvState.OBJECT_TYPE_INDEX: manual_tweaks_initial_state.append( torch.zeros([len(state)])) continue changed_state = state * 1 # just to copy tweak_value_state = state_h if random.random() >= 0.5: tweak_value_state *= -1 changed_state[arg_ind] += tweak_value_state env_state_h, force_success_flag_h, force_that_applied_h = environment.init_location_and_apply_force( initial_state=EnvState.fromTensor(changed_state), forces=forces, list_of_contact_points=list_of_contact_points) manual_tweaks_initial_state.append( (env_state_h.toTensor() - f_x_env_state_tensor) / (tweak_value_state)) finite_diff_success_flags.append(force_success_flag_h) ctx.initial_state_plus_h_diff = torch.stack( manual_tweaks_initial_state, dim=-2) finite_diff_success_flags += [force_success_flag] finite_diff_success_flags = torch.Tensor(finite_diff_success_flags) force_that_applied = torch.stack(force_that_applied, dim=0) env_state_tensor = env_state.toTensor() device = forces_tensor.device env_state_tensor = env_state_tensor.to(device=device) if not no_grad_on: ctx.f_x_plus_h_diff = ctx.f_x_plus_h_diff.to(device=device) ctx.initial_state_plus_h_diff = ctx.initial_state_plus_h_diff.to( device=device) ctx.contact_pt_x_plus_h_diff = ctx.contact_pt_x_plus_h_diff.to( device=device) return env_state_tensor, finite_diff_success_flags, force_that_applied
def forward(ctx, environment, initial_state, forces_tensor, list_of_contact_points): assert len(forces_tensor.shape) == 1, 'The force should be flattened' initial_state = EnvState.fromTensor(initial_state) forces = ForceValOnly.fromForceArray( forces_tensor.view(environment.number_of_cp, 3)) list_of_contact_points = list_of_contact_points.view(5, 3) # print("With batch Input:", initial_state, forces_tensor, list_of_contact_points) batch_phy_input = [{ 'forces': [force.tolist() for force in forces], 'initial_state': initial_state.to_dict(), 'object_num': None, 'list_of_contact_points': list_of_contact_points.cpu().tolist() }] batch_tweak_values = [None] no_grad_on = not (True in ctx.needs_input_grad) with torch.no_grad( ): # manually compute the gradients via finite difference cp_h = 0.05 contact_point_tensor = list_of_contact_points.view(-1) # to compute dS/dC for arg_ind in range(len(contact_point_tensor)): changed_contact_point = contact_point_tensor * 1 # Just copy tweak_value_cp = cp_h if random.random() > 0.5: tweak_value_cp *= -1 changed_contact_point[arg_ind] += tweak_value_cp changed_contact_point = changed_contact_point.view(5, 3) batch_tweak_values.append(tweak_value_cp) batch_phy_input.append({ 'forces': [force.tolist() for force in forces], 'initial_state': initial_state.to_dict(), 'object_num': None, 'list_of_contact_points': changed_contact_point.cpu().tolist() }) # to compute dS/dF force_h = environment.force_h for arg_ind in range(len(forces_tensor)): changed_force = forces_tensor * 1 # Just copy tweak_value_force = force_h if random.random() > 0.5: tweak_value_force *= -1 changed_force[arg_ind] += tweak_value_force changed_force = ForceValOnly.fromForceArray( changed_force.view(environment.number_of_cp, -1)) batch_tweak_values.append(tweak_value_force) batch_phy_input.append({ 'forces': [force.tolist() for force in changed_force], 'initial_state': initial_state.to_dict(), 'object_num': None, 'list_of_contact_points': list_of_contact_points.cpu().tolist() }) # to compute dS_next/dS_current state = initial_state.toTensor() state_h = environment.state_h assert EnvState.OBJECT_TYPE_INDEX == len( state) - 1, "we asssume the last idx is object type" for arg_ind in range(len(state) - 1): changed_state = state * 1 # just to copy tweak_value_state = state_h if random.random() >= 0.5: tweak_value_state *= -1 changed_state[arg_ind] += tweak_value_state batch_tweak_values.append(tweak_value_state) batch_phy_input.append({ 'forces': [force.tolist() for force in forces], 'initial_state': EnvState.fromTensor(changed_state).to_dict(), 'object_num': None, 'list_of_contact_points': list_of_contact_points.cpu().tolist() }) batch_data = environment.batch_init_locations_and_apply_force( batch_data=batch_phy_input) batch_state_tensors = [ build_env_state_from_dict(one_d['state']).toTensor() for one_d in batch_data ] batch_succ_flags = [one_d['succ'] for one_d in batch_data] batch_force_locations = [one_d['loc'] for one_d in batch_data] expected_length = 1 + len(contact_point_tensor) + len( forces_tensor) + len(state) - 1 assert len(batch_state_tensors) == expected_length, \ "Result tensor dimension is unexpected, expected %d, but get %d" % (expected_length, len(batch_state_tensors)) f_x_env_state_tensor = batch_state_tensors[0] indexes = [(1, 1 + len(contact_point_tensor)), (1 + len(contact_point_tensor), 1 + len(contact_point_tensor) + len(forces_tensor)), (1 + len(contact_point_tensor) + len(forces_tensor), 1 + len(contact_point_tensor) + len(forces_tensor) + len(state) - 1)] cp_tweak_state_tensors, force_tweak_state_tensors, next_state_tweak_tensors = \ [batch_state_tensors[a: b] for a, b in indexes] cp_tweak_values, force_tweak_values, next_state_tweak_values = \ [batch_tweak_values[a: b] for a, b in indexes] if no_grad_on: finite_diff_success_flags = [batch_succ_flags[0]] else: finite_diff_success_flags = batch_succ_flags[ 1 + len(contact_point_tensor):] + [batch_succ_flags[0]] cp_fd = [ (fx_delta - f_x_env_state_tensor) / delta for (fx_delta, delta) in zip(cp_tweak_state_tensors, cp_tweak_values) ] force_fd = [(fx_delta - f_x_env_state_tensor) / delta for ( fx_delta, delta) in zip(force_tweak_state_tensors, force_tweak_values)] state_fd = [(fx_delta - f_x_env_state_tensor) / delta for (fx_delta, delta) in zip(next_state_tweak_tensors, next_state_tweak_values)] + \ [torch.zeros([len(state)])] if not no_grad_on: ctx.contact_pt_x_plus_h_diff = torch.stack(cp_fd, dim=-2) ctx.f_x_plus_h_diff = torch.stack(force_fd, dim=-2) ctx.initial_state_plus_h_diff = torch.stack(state_fd, dim=-2) # summarize results device = forces_tensor.device force_that_applied = torch.Tensor( batch_force_locations[0]).to(device=device) finite_diff_success_flags = torch.Tensor(finite_diff_success_flags).to( device=device) f_x_env_state_tensor = f_x_env_state_tensor.to(device=device) if not no_grad_on: # if it requires grads. ctx.contact_pt_x_plus_h_diff = ctx.contact_pt_x_plus_h_diff.to( device=device) ctx.f_x_plus_h_diff = ctx.f_x_plus_h_diff.to(device=device) ctx.initial_state_plus_h_diff = ctx.initial_state_plus_h_diff.to( device=device) # print("With batch output:", f_x_env_state_tensor, force_that_applied) return f_x_env_state_tensor, finite_diff_success_flags, force_that_applied